muse_vqganvae_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import datasets\n",
    "from torchvision.io import read_image\n",
    "from torchvision.transforms import ToTensor,Compose\n",
    "from accelerate.utils import LoggerType\n",
    "\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "from torchvision.io import read_image\n",
    "from PIL import Image\n",
    "\n",
    "class CustomImageDataset(Dataset):\n",
    "\n",
    "\n",
    "    def __init__(self,  img_dir):\n",
    "        self.img_dir = img_dir\n",
    "\n",
    "    def __len__(self):\n",
    "        if(len(self.all_files) == 0):\n",
    "            all_files = []\n",
    "            for root, dirs, files in sorted(os.walk(self.img_dir)):\n",
    "                for file in files:\n",
    "                    all_files.append(root + '/' + file)\n",
    "                    self.all_files = all_files\n",
    "            print(str(len(self.all_files)) + \" images found.\")\n",
    "        return len(self.all_files)\n",
    "\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        trans = Compose([ToTensor()])\n",
    "        if(len(self.all_files) == 0):\n",
    "            all_files = []\n",
    "            for root, dirs, files in sorted(os.walk(self.img_dir)):\n",
    "                for file in files:\n",
    "                    all_files.append(root + '/' + file)\n",
    "                    self.all_files = all_files\n",
    "        image = Image.open(self.all_files[idx]).convert('RGB')\n",
    "        image = trans(image)\n",
    "        \n",
    "        return image\n",
    "\n",
    "training_data = CustomImageDataset('/media/gamal/Passport/Datasets/VoxCeleb2Test/Voxceleb2TestFaces')\n",
    "\n",
    "\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "\n",
    "accelerator = Accelerator(log_with=\"tensorboard\",\n",
    "                          project_dir='/media/gamal/Passport/muse/vqganvae_log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from muse_maskgit_pytorch import VQGanVAE, VQGanVAETrainer\n",
    "\n",
    "import os\n",
    "import configparser\n",
    "\n",
    "# Loading configurations\n",
    "configParser = configparser.RawConfigParser()   \n",
    "configFilePath = r'configuration.txt'\n",
    "configParser.read(configFilePath)\n",
    "datasetPathVideo =  configParser.get('COMMON', 'datasetPathVideo')\n",
    "\n",
    "\n",
    "vae = VQGanVAE(\n",
    "    dim = 32,\n",
    "    vq_codebook_dim = 8192,\n",
    "    vq_codebook_size = 8192,\n",
    "    channels=3,\n",
    "    \n",
    ")\n",
    "\n",
    "vae.load('/media/gamal/Passport/muse/vqganvae/vae.12145000.pt')\n",
    "# train on folder of images, as many images as possible\n",
    "\n",
    "\n",
    "trainer = VQGanVAETrainer(\n",
    "    current_step=1,\n",
    "    num_epochs=50000,\n",
    "    vae = vae,\n",
    "    dataloader=train_dataloader,\n",
    "    valid_dataloader=test_dataloader,\n",
    "    gradient_accumulation_steps = 1,\n",
    "    num_train_steps = 50000,\n",
    "    results_dir = '/media/gamal/Passport/muse/vqganvae',\n",
    "    save_results_every = 100,\n",
    "    save_model_every = 100000,\n",
    "    only_save_last_checkpoint=True,\n",
    "    accelerator=accelerator,\n",
    ").cuda()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "trainer.train()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ds2f_m_i",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}