model traine upaded 2.O

lightning_logs/version_20/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_21/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_22/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_23/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_24/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_25/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_26/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_27/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_28/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_29/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_30/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_31/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_32/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_33/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 3

lightning_logs/version_34/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_35/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_36/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

params.yaml

@@ -28,6 +28,4 @@ IN_CHANNELS: 3
 LEARNING_RATE : 2e-4
 LAMBDA_RECON : 100
 DISPLAY_STEP : 10
-INPUT_CHANNELS : 1
-OUTPUT_CHANNELS : 2
 EPOCH : 1

research/data_transformation.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -21,7 +21,7 @@
        "'c:\\\\mlops project\\\\image-colorization-mlops'"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,7 +88,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -97,27 +97,31 @@
     "from torch.utils.data import Dataset\n",
     "from torchvision import transforms\n",
     "\n",
-    "class ImageColorizationDataset(Dataset):\n",
-    "    def __init__(self, dataset, image_size, transform=None):\n",
+    "        \n",
+    "        \n",
+    "        \n",
+    "class ImageColorizationDataset:\n",
+    "    def __init__(self, dataset, image_size, transform = None):\n",
     "        self.dataset = dataset\n",
     "        self.transform = transform\n",
-    "        self.image_size = tuple(image_size)  \n",
+    "        self.image_size = tuple(image_size)\n",
+    "        \n",
     "    def __len__(self):\n",
     "        return len(self.dataset[0])\n",
     "    \n",
     "    def __getitem__(self, idx):\n",
-    "        L = np.array(self.dataset[0][idx]).reshape(self.image_size)\n",
+    "        L  = np.array(self.dataset[0][idx]).reshape(self.image_size)\n",
     "        L = transforms.ToTensor()(L)\n",
     "        \n",
     "        ab = np.array(self.dataset[1][idx])\n",
     "        ab = transforms.ToTensor()(ab)\n",
     "        \n",
-    "        return ab, L"
+    "        return ab , L"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -142,9 +146,11 @@
     "    def get_datasets(self, dataset):\n",
     "        train_dataset = ImageColorizationDataset(\n",
     "            dataset=dataset,\n",
+    "            image_size=self.config.IMAGE_SIZE\n",
     "        )\n",
     "        test_dataset = ImageColorizationDataset(\n",
     "            dataset=dataset,\n",
+    "            image_size=self.config.IMAGE_SIZE\n",
     "        )\n",
     "    \n",
     "        return train_dataset, test_dataset\n",
@@ -172,18 +178,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[2024-08-24 19:09:25,021: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
-      "[2024-08-24 19:09:25,024: INFO: common: yaml file: params.yaml loaded successfully]\n",
-      "[2024-08-24 19:09:25,026: INFO: common: created directory at: artifacts]\n",
-      "[2024-08-24 19:09:43,417: INFO: 3400243030: Train dataset saved at: artifacts/data_transformation\\train_dataset.pt]\n",
-      "[2024-08-24 19:09:43,440: INFO: 3400243030: Test dataset saved at: artifacts/data_transformation\\test_dataset.pt]\n"
+      "[2024-08-24 23:35:18,235: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
+      "[2024-08-24 23:35:18,243: INFO: common: yaml file: params.yaml loaded successfully]\n",
+      "[2024-08-24 23:35:18,245: INFO: common: created directory at: artifacts]\n",
+      "[2024-08-24 23:35:34,541: INFO: 2411080742: Train dataset saved at: artifacts/data_transformation\\train_dataset.pt]\n",
+      "[2024-08-24 23:35:34,552: INFO: 2411080742: Test dataset saved at: artifacts/data_transformation\\test_dataset.pt]\n",
+      "Train dataset type: <class '__main__.ImageColorizationDataset'>\n",
+      "Train dataset length: 5000\n",
+      "First item in train dataset: (tensor([[[0.5059, 0.4941, 0.4941,  ..., 0.4902, 0.4863, 0.4863],\n",
+      "         [0.4980, 0.4941, 0.4941,  ..., 0.4980, 0.4902, 0.4824],\n",
+      "         [0.4980, 0.4980, 0.4980,  ..., 0.4902, 0.4941, 0.5020],\n",
+      "         ...,\n",
+      "         [0.4941, 0.4941, 0.4980,  ..., 0.4941, 0.4941, 0.4941],\n",
+      "         [0.4941, 0.4980, 0.4941,  ..., 0.4941, 0.4941, 0.4941],\n",
+      "         [0.4980, 0.4980, 0.5020,  ..., 0.4941, 0.4863, 0.4941]],\n",
+      "\n",
+      "        [[0.5333, 0.5255, 0.5294,  ..., 0.5176, 0.5255, 0.5294],\n",
+      "         [0.5373, 0.5333, 0.5216,  ..., 0.5137, 0.5216, 0.5373],\n",
+      "         [0.5412, 0.5255, 0.5255,  ..., 0.5137, 0.5137, 0.5216],\n",
+      "         ...,\n",
+      "         [0.5137, 0.5137, 0.5137,  ..., 0.5098, 0.5098, 0.5098],\n",
+      "         [0.5137, 0.5137, 0.5176,  ..., 0.5098, 0.5098, 0.5098],\n",
+      "         [0.5137, 0.5216, 0.5294,  ..., 0.5098, 0.5098, 0.5098]]]), tensor([[[0.9294, 0.5294, 0.2941,  ..., 0.1373, 0.1451, 0.2471],\n",
+      "         [0.9176, 0.5961, 0.2824,  ..., 0.1608, 0.1922, 0.1843],\n",
+      "         [0.8431, 0.8471, 0.4078,  ..., 0.2863, 0.1882, 0.3216],\n",
+      "         ...,\n",
+      "         [0.1569, 0.1765, 0.1490,  ..., 0.0431, 0.0314, 0.0314],\n",
+      "         [0.1569, 0.2196, 0.1843,  ..., 0.0314, 0.0275, 0.0392],\n",
+      "         [0.1647, 0.2353, 0.3098,  ..., 0.0471, 0.0510, 0.0588]]]))\n"
      ]
     }
    ],
@@ -196,22 +225,350 @@
     "    # Load the dataset\n",
     "    dataset = data_transformation.load_data()\n",
     "    \n",
-    "    # Get the dataloader using the loaded dataset\n",
+    "    # Get the datasets using the loaded dataset\n",
     "    train_dataset, test_dataset = data_transformation.get_datasets(dataset)\n",
     "    \n",
     "    # Perform any further operations (e.g., saving the dataset)\n",
     "    data_transformation.save_datasets(train_dataset, test_dataset)\n",
     "    \n",
+    "    # Print information about the train dataset\n",
+    "    print(f\"Train dataset type: {type(train_dataset)}\")\n",
+    "    print(f\"Train dataset length: {len(train_dataset)}\")\n",
+    "    print(f\"First item in train dataset: {train_dataset[0]}\")\n",
+    "    \n",
     "except Exception as e:\n",
-    "    raise e\n"
+    "    raise e"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\azizu\\AppData\\Local\\Temp\\ipykernel_28412\\4121028732.py:4: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  test_dataset = torch.load('artifacts/data_transformation/test_dataset.pt')\n",
+      "C:\\Users\\azizu\\AppData\\Local\\Temp\\ipykernel_28412\\4121028732.py:5: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  train_dataset = torch.load('artifacts/data_transformation/train_dataset.pt')\n"
+     ]
+    }
+   ],
+   "source": [
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "\n",
+    "\n",
+    "test_dataset = torch.load('artifacts/data_transformation/test_dataset.pt')\n",
+    "train_dataset = torch.load('artifacts/data_transformation/train_dataset.pt')\n",
+    "\n",
+    "\n",
+    "train_loader1 = DataLoader(\n",
+    "    train_dataset,\n",
+    "    shuffle=True,\n",
+    "    batch_size=1\n",
+    ")\n",
+    "test_loader1 = DataLoader(\n",
+    "    test_dataset,\n",
+    "    shuffle=True,\n",
+    "    batch_size=1\n",
+    ")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train loader batch - ab shape: torch.Size([1, 2, 224, 224]), L shape: torch.Size([1, 1, 224, 224])\n",
+      "Test loader batch - ab shape: torch.Size([1, 2, 224, 224]), L shape: torch.Size([1, 1, 224, 224])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Print shapes for train dataloader\n",
+    "for batch in train_loader1:\n",
+    "    ab, L = batch\n",
+    "    print(f\"Train loader batch - ab shape: {ab.shape}, L shape: {L.shape}\")\n",
+    "    break  # We only need to check one batch\n",
+    "\n",
+    "# Print shapes for test dataloader\n",
+    "for batch in test_loader1:\n",
+    "     ab, L = batch\n",
+    "     print(f\"Test loader batch - ab shape: {ab.shape}, L shape: {L.shape}\")\n",
+    "     break  # We only need to check one batch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "from skimage.color import rgb2lab, lab2rgb\n",
+    "import gc\n",
+    "import matplotlib.pylab as plt\n",
+    "\n",
+    "def lab_to_rgb(L, ab):\n",
+    "    L = L * 100\n",
+    "    ab = (ab - 0.5) * 128 * 2\n",
+    "    Lab = torch.cat([L, ab], dim = 2).numpy()\n",
+    "    rgb_img = []\n",
+    "    for img in Lab:\n",
+    "        img_rgb = lab2rgb(img)\n",
+    "        rgb_img.append(img_rgb)\n",
+    "        \n",
+    "    return np.stack(rgb_img, axis = 0)\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def display_progress(cond, real, fake, current_epoch = 0, figsize=(20,15)):\n",
+    "    \"\"\"\n",
+    "    Save cond, real (original) and generated (fake)\n",
+    "    images in one panel \n",
+    "    \"\"\"\n",
+    "    cond = cond.detach().cpu().permute(1, 2, 0)   \n",
+    "    real = real.detach().cpu().permute(1, 2, 0)\n",
+    "    fake = fake.detach().cpu().permute(1, 2, 0)\n",
+    "    \n",
+    "    images = [cond, real, fake]\n",
+    "    titles = ['input','real','generated']\n",
+    "    print(f'Epoch: {current_epoch}')\n",
+    "    fig, ax = plt.subplots(1, 3, figsize=figsize)\n",
+    "    for idx,img in enumerate(images):\n",
+    "        if idx == 0:\n",
+    "            ab = torch.zeros((224,224,2))\n",
+    "            img = torch.cat([images[0]* 100, ab], dim=2).numpy()\n",
+    "            imgan = lab2rgb(img)\n",
+    "        else:\n",
+    "            imgan = lab_to_rgb(images[0],img)\n",
+    "        ax[idx].imshow(imgan)\n",
+    "        ax[idx].axis(\"off\")\n",
+    "    for idx, title in enumerate(titles):    \n",
+    "        ax[idx].set_title('{}'.format(title))\n",
+    "    plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pytorch_lightning as pl\n",
+    "from src.imagecolorization.conponents.model_trainer import Generator,Critic\n",
+    "from torch import nn, optim\n",
+    "\n",
+    "class CWGAN(pl.LightningModule):\n",
+    "    def __init__(self, in_channels, out_channels, learning_rate=0.0002, lambda_recon=100, display_step=10, lambda_gp=10, lambda_r1=10):\n",
+    "        super().__init__()\n",
+    "        self.save_hyperparameters()\n",
+    "        self.display_step = display_step\n",
+    "        self.generator = Generator(in_channels, out_channels)\n",
+    "        self.critic = Critic(in_channels + out_channels)  # Ensure Critic is initialized with the correct input channels\n",
+    "        self.lambda_recon = lambda_recon\n",
+    "        self.lambda_gp = lambda_gp\n",
+    "        self.lambda_r1 = lambda_r1\n",
+    "        self.recon_criterion = nn.L1Loss()\n",
+    "        self.generator_losses, self.critic_losses = [], []\n",
+    "        self.automatic_optimization = False\n",
+    "        \n",
+    "    def configure_optimizers(self):\n",
+    "        optimizer_G = optim.Adam(self.generator.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.9))\n",
+    "        optimizer_C = optim.Adam(self.critic.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.9))\n",
+    "        return [optimizer_C, optimizer_G]\n",
+    "    \n",
+    "    def generator_step(self, real_images, conditioned_images, optimizer_G):\n",
+    "        # WGAN has only a reconstruction loss\n",
+    "        optimizer_G.zero_grad()\n",
+    "        fake_images = self.generator(conditioned_images)\n",
+    "        recon_loss = self.recon_criterion(fake_images, real_images)\n",
+    "        recon_loss.backward()\n",
+    "        optimizer_G.step()\n",
+    "        self.generator_losses.append(recon_loss.item())\n",
+    "        \n",
+    "    def critic_step(self, real_images, conditioned_images, optimizer_C):\n",
+    "        optimizer_C.zero_grad()\n",
+    "        fake_images = self.generator(conditioned_images)\n",
+    "        fake_input = torch.cat((fake_images, conditioned_images), 1)\n",
+    "        real_input = torch.cat((real_images, conditioned_images), 1)\n",
+    "        fake_logits = self.critic(fake_input)\n",
+    "        real_logits = self.critic(real_input)\n",
+    "\n",
+    "        # Compute the loss for the critic\n",
+    "        loss_C = real_logits.mean() - fake_logits.mean()\n",
+    "\n",
+    "        # Compute the gradient penalty\n",
+    "        alpha = torch.rand(real_images.size(0), 1, 1, 1, requires_grad=True).to(real_images.device)\n",
+    "        interpolated = (alpha * real_images + (1 - alpha) * fake_images.detach()).requires_grad_(True)\n",
+    "        interpolated_logits = self.critic(interpolated, conditioned_images)\n",
+    "    \n",
+    "        gradients = torch.autograd.grad(outputs=interpolated_logits, inputs=interpolated,\n",
+    "                                    grad_outputs=torch.ones_like(interpolated_logits), create_graph=True, retain_graph=True)[0]\n",
+    "        gradients = gradients.view(len(gradients), -1)\n",
+    "        gradients_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n",
+    "        loss_C += self.lambda_gp * gradients_penalty\n",
+    "    \n",
+    "        # Compute the R1 regularization loss\n",
+    "        r1_reg = gradients.pow(2).sum(1).mean()\n",
+    "        loss_C += self.lambda_r1 * r1_reg\n",
+    "\n",
+    "        # Backpropagation\n",
+    "        loss_C.backward()\n",
+    "        optimizer_C.step()\n",
+    "        self.critic_losses.append(loss_C.item())\n",
+    "\n",
+    "    \n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        real, condition = batch\n",
+    "        optimizer_C, optimizer_G = self.optimizers()  # Access optimizers\n",
+    "    \n",
+    "        # Debugging shapes\n",
+    "        print(f\"Real images shape: {real.shape}\")\n",
+    "        print(f\"Conditioned images shape: {condition.shape}\")\n",
+    "    \n",
+    "        # Update the critic\n",
+    "        self.critic_step(real, condition, optimizer_C)\n",
+    "    \n",
+    "        # Update the generator\n",
+    "        self.generator_step(real, condition, optimizer_G)\n",
+    "    \n",
+    "    \n",
+    "        # Logging and saving models\n",
+    "        gen_mean = sum(self.generator_losses[-self.display_step:]) / self.display_step\n",
+    "        crit_mean = sum(self.critic_losses[-self.display_step:]) / self.display_step\n",
+    "        if self.current_epoch % self.display_step == 0 and batch_idx == 0:\n",
+    "            fake = self.generator(condition).detach()\n",
+    "            print(f\"Epoch {self.current_epoch}: Generator loss: {gen_mean}, Critic loss: {crit_mean}\")\n",
+    "            display_progress(condition[0], real[0], fake[0], self.current_epoch)\n",
+    "    \n",
+    "        # Save models every 10 epochs\n",
+    "        if (self.current_epoch + 1) % 10 == 0 and batch_idx == self.trainer.num_training_batches - 1:\n",
+    "            torch.save(self.generator.state_dict(), f\"/kaggle/working/cwgan_generator_epoch_{self.current_epoch+1}.pt\")\n",
+    "            torch.save(self.critic.state_dict(), f\"/kaggle/working/cwgan_critic_epoch_{self.current_epoch+1}.pt\")\n",
+    "            print(f\"Saved models at epoch {self.current_epoch+1}\")\n",
+    "\n",
+    "        # Final save at epoch 150\n",
+    "        if self.current_epoch == 149 and batch_idx == self.trainer.num_training_batches - 1:\n",
+    "            torch.save(self.generator.state_dict(), \"cwgan_generator_final.pt\")\n",
+    "            torch.save(self.critic.state_dict(), \"cwgan_critic_final.pt\")\n",
+    "            print(\"Saved final models at epoch 150\")\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch import nn, optim\n",
+    "\n",
+    "gc.collect()\n",
+    "cwgan = CWGAN(in_channels = 1, out_channels = 2 ,learning_rate=2e-4, lambda_recon=100, display_step=10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2024-08-25 00:04:53,828: INFO: rank_zero: GPU available: True (cuda), used: True]\n",
+      "[2024-08-25 00:04:53,829: INFO: rank_zero: TPU available: False, using: 0 TPU cores]\n",
+      "[2024-08-25 00:04:53,830: INFO: rank_zero: IPU available: False, using: 0 IPUs]\n",
+      "[2024-08-25 00:04:53,831: INFO: rank_zero: HPU available: False, using: 0 HPUs]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2024-08-25 00:04:54,549: INFO: cuda: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]]\n",
+      "[2024-08-25 00:04:54,552: INFO: model_summary: \n",
+      "  | Name            | Type      | Params\n",
+      "----------------------------------------------\n",
+      "0 | generator       | Generator | 8.2 M \n",
+      "1 | critic          | Critic    | 2.8 M \n",
+      "2 | recon_criterion | L1Loss    | 0     \n",
+      "----------------------------------------------\n",
+      "11.0 M    Trainable params\n",
+      "0         Non-trainable params\n",
+      "11.0 M    Total params\n",
+      "43.893    Total estimated model params size (MB)]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a1c97e80cdd746b2afa3ae30771ee058",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Training: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Real images shape: torch.Size([1, 2, 224, 224])\n",
+      "Conditioned images shape: torch.Size([1, 1, 224, 224])\n"
+     ]
+    },
+    {
+     "ename": "TypeError",
+     "evalue": "Critic.forward() missing 1 required positional argument: 'l'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[1;32mIn[15], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m trainer \u001b[38;5;241m=\u001b[39m pl\u001b[38;5;241m.\u001b[39mTrainer(max_epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m150\u001b[39m)\n\u001b[1;32m----> 2\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcwgan\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader1\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\trainer.py:545\u001b[0m, in \u001b[0;36mTrainer.fit\u001b[1;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[0;32m    543\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mstatus \u001b[38;5;241m=\u001b[39m TrainerStatus\u001b[38;5;241m.\u001b[39mRUNNING\n\u001b[0;32m    544\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m--> 545\u001b[0m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_and_handle_interrupt\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    546\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fit_impl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_dataloaders\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdatamodule\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\n\u001b[0;32m    547\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\call.py:44\u001b[0m, in \u001b[0;36m_call_and_handle_interrupt\u001b[1;34m(trainer, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[0;32m     42\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m     43\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher\u001b[38;5;241m.\u001b[39mlaunch(trainer_fn, \u001b[38;5;241m*\u001b[39margs, trainer\u001b[38;5;241m=\u001b[39mtrainer, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m---> 44\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrainer_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     46\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _TunerExitException:\n\u001b[0;32m     47\u001b[0m     _call_teardown_hook(trainer)\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\trainer.py:581\u001b[0m, in \u001b[0;36mTrainer._fit_impl\u001b[1;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[0;32m    574\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mfn \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m    575\u001b[0m ckpt_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_checkpoint_connector\u001b[38;5;241m.\u001b[39m_select_ckpt_path(\n\u001b[0;32m    576\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mfn,\n\u001b[0;32m    577\u001b[0m     ckpt_path,\n\u001b[0;32m    578\u001b[0m     model_provided\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[0;32m    579\u001b[0m     model_connected\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m    580\u001b[0m )\n\u001b[1;32m--> 581\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mckpt_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mckpt_path\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    583\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mstopped\n\u001b[0;32m    584\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\trainer.py:990\u001b[0m, in \u001b[0;36mTrainer._run\u001b[1;34m(self, model, ckpt_path)\u001b[0m\n\u001b[0;32m    985\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_signal_connector\u001b[38;5;241m.\u001b[39mregister_signal_handlers()\n\u001b[0;32m    987\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[0;32m    988\u001b[0m \u001b[38;5;66;03m# RUN THE TRAINER\u001b[39;00m\n\u001b[0;32m    989\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[1;32m--> 990\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_stage\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    992\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[0;32m    993\u001b[0m \u001b[38;5;66;03m# POST-Training CLEAN UP\u001b[39;00m\n\u001b[0;32m    994\u001b[0m \u001b[38;5;66;03m# ----------------------------\u001b[39;00m\n\u001b[0;32m    995\u001b[0m log\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m: trainer tearing down\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\trainer.py:1036\u001b[0m, in \u001b[0;36mTrainer._run_stage\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m   1034\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_sanity_check()\n\u001b[0;32m   1035\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mset_detect_anomaly(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_detect_anomaly):\n\u001b[1;32m-> 1036\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1037\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   1038\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUnexpected state \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\fit_loop.py:202\u001b[0m, in \u001b[0;36m_FitLoop.run\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    200\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m    201\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_start()\n\u001b[1;32m--> 202\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    203\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end()\n\u001b[0;32m    204\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\fit_loop.py:359\u001b[0m, in \u001b[0;36m_FitLoop.advance\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    357\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_training_epoch\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m    358\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_data_fetcher \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m--> 359\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mepoch_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_data_fetcher\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\training_epoch_loop.py:136\u001b[0m, in \u001b[0;36m_TrainingEpochLoop.run\u001b[1;34m(self, data_fetcher)\u001b[0m\n\u001b[0;32m    134\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdone:\n\u001b[0;32m    135\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 136\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata_fetcher\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    137\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_advance_end(data_fetcher)\n\u001b[0;32m    138\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\training_epoch_loop.py:242\u001b[0m, in \u001b[0;36m_TrainingEpochLoop.advance\u001b[1;34m(self, data_fetcher)\u001b[0m\n\u001b[0;32m    240\u001b[0m             batch_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mautomatic_optimization\u001b[38;5;241m.\u001b[39mrun(trainer\u001b[38;5;241m.\u001b[39moptimizers[\u001b[38;5;241m0\u001b[39m], batch_idx, kwargs)\n\u001b[0;32m    241\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 242\u001b[0m             batch_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmanual_optimization\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    244\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbatch_progress\u001b[38;5;241m.\u001b[39mincrement_processed()\n\u001b[0;32m    246\u001b[0m \u001b[38;5;66;03m# update non-plateau LR schedulers\u001b[39;00m\n\u001b[0;32m    247\u001b[0m \u001b[38;5;66;03m# update epoch-interval ones only when we are at the end of training epoch\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\optimization\\manual.py:92\u001b[0m, in \u001b[0;36m_ManualOptimization.run\u001b[1;34m(self, kwargs)\u001b[0m\n\u001b[0;32m     90\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_run_start()\n\u001b[0;32m     91\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m suppress(\u001b[38;5;167;01mStopIteration\u001b[39;00m):  \u001b[38;5;66;03m# no loop to break at this level\u001b[39;00m\n\u001b[1;32m---> 92\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madvance\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     93\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_restarting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m     94\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mon_run_end()\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\loops\\optimization\\manual.py:112\u001b[0m, in \u001b[0;36m_ManualOptimization.advance\u001b[1;34m(self, kwargs)\u001b[0m\n\u001b[0;32m    109\u001b[0m trainer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\n\u001b[0;32m    111\u001b[0m \u001b[38;5;66;03m# manually capture logged metrics\u001b[39;00m\n\u001b[1;32m--> 112\u001b[0m training_step_output \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_strategy_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrainer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtraining_step\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalues\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    113\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m kwargs  \u001b[38;5;66;03m# release the batch from memory\u001b[39;00m\n\u001b[0;32m    114\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mpost_training_step()  \u001b[38;5;66;03m# unused hook - call anyway for backward compatibility\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\trainer\\call.py:309\u001b[0m, in \u001b[0;36m_call_strategy_hook\u001b[1;34m(trainer, hook_name, *args, **kwargs)\u001b[0m\n\u001b[0;32m    306\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m    308\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m trainer\u001b[38;5;241m.\u001b[39mprofiler\u001b[38;5;241m.\u001b[39mprofile(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[Strategy]\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtrainer\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mhook_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m--> 309\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    311\u001b[0m \u001b[38;5;66;03m# restore current_fx when nested context\u001b[39;00m\n\u001b[0;32m    312\u001b[0m pl_module\u001b[38;5;241m.\u001b[39m_current_fx_name \u001b[38;5;241m=\u001b[39m prev_fx_name\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\pytorch_lightning\\strategies\\strategy.py:382\u001b[0m, in \u001b[0;36mStrategy.training_step\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    380\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module:\n\u001b[0;32m    381\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_redirection(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_module, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtraining_step\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m--> 382\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlightning_module\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining_step\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[1;32mIn[13], line 74\u001b[0m, in \u001b[0;36mCWGAN.training_step\u001b[1;34m(self, batch, batch_idx)\u001b[0m\n\u001b[0;32m     71\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mConditioned images shape: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcondition\u001b[38;5;241m.\u001b[39mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     73\u001b[0m \u001b[38;5;66;03m# Update the critic\u001b[39;00m\n\u001b[1;32m---> 74\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcritic_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreal\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer_C\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     76\u001b[0m \u001b[38;5;66;03m# Update the generator\u001b[39;00m\n\u001b[0;32m     77\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgenerator_step(real, condition, optimizer_G)\n",
+      "Cell \u001b[1;32mIn[13], line 38\u001b[0m, in \u001b[0;36mCWGAN.critic_step\u001b[1;34m(self, real_images, conditioned_images, optimizer_C)\u001b[0m\n\u001b[0;32m     36\u001b[0m fake_input \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mcat((fake_images, conditioned_images), \u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m     37\u001b[0m real_input \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mcat((real_images, conditioned_images), \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m---> 38\u001b[0m fake_logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcritic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfake_input\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     39\u001b[0m real_logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcritic(real_input)\n\u001b[0;32m     41\u001b[0m \u001b[38;5;66;03m# Compute the loss for the critic\u001b[39;00m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "\u001b[1;31mTypeError\u001b[0m: Critic.forward() missing 1 required positional argument: 'l'"
+     ]
+    }
+   ],
+   "source": [
+    "trainer = pl.Trainer(max_epochs=150)\n",
+    "trainer.fit(cwgan, train_loader1)"
+   ]
   },
   {
    "cell_type": "code",

research/lightning_logs/version_0/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

research/lightning_logs/version_1/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

research/model_building.ipynb

@@ -183,7 +183,7 @@
     "    \n",
     "    \n",
     "class Critic(nn.Module):\n",
-    "    def __init__(self, in_channels=3):\n",
+    "    def __init__(self, in_channels):\n",
     "        super(Critic, self).__init__()\n",
     "\n",
     "        def critic_block(in_filters, out_filters, normalization=True):\n",
@@ -203,11 +203,9 @@
     "            nn.Linear(512, 1)\n",
     "        )\n",
     "\n",
-    "    def forward(self, ab, l):\n",
-    "        img_input = torch.cat((ab, l), 1)\n",
+    "    def forward(self, img_input):\n",
     "        output = self.model(img_input)\n",
     "        return output\n",
-    "    \n",
     "    "
    ]
   },

src/imagecolorization/config/configuration.py

@@ -91,6 +91,7 @@ class ConfigurationManager:
             BATCH_SIZE= params.BATCH_SIZE
         )
         return model_trainer_config
+
         
     
     

src/imagecolorization/conponents/model_building.py

@@ -4,67 +4,68 @@ from pathlib import Path
 from torchsummary import summary
 from src.imagecolorization.entity.config_entity import ModelBuildingConfig
 class ResBlock(nn.Module):
-    def __init__(self, in_channles, out_channels, stride = 1, kerenl_size = 3, padding = 1, bias = False):
+    def __init__(self, in_channels, out_channels, stride=1):
         super().__init__()
         self.layer = nn.Sequential(
-            nn.Conv2d(in_channles, out_channels, kernel_size=kerenl_size, padding=padding, stride=stride, bias = bias),
+            nn.Conv2d(in_channels,out_channels,kernel_size=3, padding=1, stride=stride, bias=False),
             nn.BatchNorm2d(out_channels),
             nn.ReLU(inplace=True),
-            nn.Conv2d(out_channels, out_channels, kernel_size=kerenl_size, padding=padding, stride = 1, bias = bias),
+            nn.Conv2d(out_channels, out_channels,kernel_size=3,padding=1, stride=1, bias=False),
             nn.BatchNorm2d(out_channels),
             nn.ReLU(inplace=True)
         )
-        
-        self.identity_map = nn.Conv2d(in_channles, out_channels,kernel_size=1, stride=stride)
-        self.relu = nn.ReLU(inplace= True)
-        
+
+        self.identity_map = nn.Conv2d(in_channels, out_channels,kernel_size=1,stride=stride)
+        self.relu = nn.ReLU(inplace=True)
     def forward(self, inputs):
         x = inputs.clone().detach()
         out = self.layer(x)
-        residual = self.identity_map(inputs)
-        skip  = out + residual
+        residual  = self.identity_map(inputs)
+        skip = out + residual
         return self.relu(skip)
     
     
-class DownsampleConv(nn.Module):
-    def __init__(self, in_channels, out_channels, stride = 1):
+class DownSampleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
         super().__init__()
         self.layer = nn.Sequential(
             nn.MaxPool2d(2),
             ResBlock(in_channels, out_channels)
         )
-        
+
     def forward(self, inputs):
         return self.layer(inputs)
     
     
     
-class UpsampleConv(nn.Module):
-    def __init__(self, in_channels, out_channels, scale_factor=2):
+class UpSampleConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
         super().__init__()
-        self.upsample = nn.Upsample(scale_factor=scale_factor,mode = 'bilinear', align_corners=True)
+        
+        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
         self.res_block = ResBlock(in_channels + out_channels, out_channels)
-
+        
     def forward(self, inputs, skip):
         x = self.upsample(inputs)
-        x = torch.cat([x, skip], dim = 1)
+        x = torch.cat([x, skip], dim=1)
         x = self.res_block(x)
         return x
     
 class Generator(nn.Module):
-    def __init__(self, input_channels, output_channels, dropout_rate = 0.2):
+    def __init__(self, input_channel, output_channel, dropout_rate = 0.2):
         super().__init__()
-        self.encoding_layer1_= ResBlock(input_channels, 64)
-        self.encoding_layer2_ = DownsampleConv(64, 128)
-        self.encoding_layer3_ = DownsampleConv(128, 256)
-        self.bridge = DownsampleConv(256, 512)
-        self.decoding_layer3 = UpsampleConv(512, 256)
-        self.decoding_layer2 = UpsampleConv(256, 128)
-        self.decoding_layer1 = UpsampleConv(128 , 64)
-        self.output = nn.Conv2d(64, output_channels, kernel_size = 1)
+        self.encoding_layer1_ = ResBlock(input_channel,64)
+        self.encoding_layer2_ = DownSampleConv(64, 128)
+        self.encoding_layer3_ = DownSampleConv(128, 256)
+        self.bridge = DownSampleConv(256, 512)
+        self.decoding_layer3_ = UpSampleConv(512, 256)
+        self.decoding_layer2_ = UpSampleConv(256, 128)
+        self.decoding_layer1_ = UpSampleConv(128, 64)
+        self.output = nn.Conv2d(64, output_channel, kernel_size=1)
         self.dropout = nn.Dropout2d(dropout_rate)
         
     def forward(self, inputs):
+        ###################### Enocoder #########################
         e1 = self.encoding_layer1_(inputs)
         e1 = self.dropout(e1)
         e2 = self.encoding_layer2_(e1)
@@ -72,22 +73,25 @@ class Generator(nn.Module):
         e3 = self.encoding_layer3_(e2)
         e3 = self.dropout(e3)
         
+        ###################### Bridge #########################
         bridge = self.bridge(e3)
         bridge = self.dropout(bridge)
         
-        d3 = self.decoding_layer3(bridge, e3)
-        d2 =self.decoding_layer2(d3, e2)
-        d1 = self.decoding_layer1(d2, e1)
+        ###################### Decoder #########################
+        d3 = self.decoding_layer3_(bridge, e3)
+        d2 = self.decoding_layer2_(d3, e2)
+        d1 = self.decoding_layer1_(d2, e1)
         
-        output = self.dropout(d1)
+        ###################### Output #########################
+        output = self.output(d1)
         return output
     
-    
 class Critic(nn.Module):
     def __init__(self, in_channels=3):
         super(Critic, self).__init__()
 
         def critic_block(in_filters, out_filters, normalization=True):
+            """Returns layers of each critic block"""
             layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
             if normalization:
                 layers.append(nn.InstanceNorm2d(out_filters))
@@ -105,12 +109,12 @@ class Critic(nn.Module):
         )
 
     def forward(self, ab, l):
+        # Concatenate image and condition image by channels to produce input
         img_input = torch.cat((ab, l), 1)
         output = self.model(img_input)
         return output
     
-    
-    from torchsummary import summary
+from torchsummary import summary
 import torch
 import os
 
@@ -127,8 +131,8 @@ class ModelBuilding:
 
     def get_generator(self):
         return Generator(
-            input_channels=self.config.INPUT_CHANNELS,  # corrected argument name
-            output_channels=self.config.OUTPUT_CHANNELS,  # corrected argument name
+            input_channel=self.config.INPUT_CHANNELS,  # corrected argument name
+            output_channel=self.config.OUTPUT_CHANNELS,  # corrected argument name
             dropout_rate=self.config.DROPOUT_RATE
         ).to(self.device)
 

src/imagecolorization/conponents/model_trainer.py

@@ -73,47 +73,41 @@ class CWGAN(pl.LightningModule):
         self.recon_criterion = nn.L1Loss()
         self.generator_losses, self.critic_losses = [], []
         self.automatic_optimization = False  # Disable automatic optimization
-
+        
     def configure_optimizers(self):
         optimizer_G = optim.Adam(self.generator.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.9))
         optimizer_C = optim.Adam(self.critic.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.9))
         return [optimizer_C, optimizer_G]
-
+    
     def generator_step(self, real_images, conditioned_images, optimizer_G):
+        # WGAN has only a reconstruction loss
         optimizer_G.zero_grad()
         fake_images = self.generator(conditioned_images)
         recon_loss = self.recon_criterion(fake_images, real_images)
         recon_loss.backward()
         optimizer_G.step()
         self.generator_losses.append(recon_loss.item())
-
+        
     def critic_step(self, real_images, conditioned_images, optimizer_C):
         optimizer_C.zero_grad()
         fake_images = self.generator(conditioned_images)
-
-        # Separate L and ab channels
-        l_real = conditioned_images
-        ab_real = real_images
-        ab_fake = fake_images
-
-        # Compute logits
-        fake_logits = self.critic(ab_fake, l_real)  # Pass two arguments
-        real_logits = self.critic(ab_real, l_real)  # Pass two arguments
-
+        fake_logits = self.critic(fake_images, conditioned_images)
+        real_logits = self.critic(real_images, conditioned_images)
+        
         # Compute the loss for the critic
         loss_C = real_logits.mean() - fake_logits.mean()
 
         # Compute the gradient penalty
         alpha = torch.rand(real_images.size(0), 1, 1, 1, requires_grad=True).to(real_images.device)
-        interpolated = (alpha * ab_real + (1 - alpha) * ab_fake.detach()).requires_grad_(True)
-        interpolated_logits = self.critic(interpolated, l_real)
-
+        interpolated = (alpha * real_images + (1 - alpha) * fake_images.detach()).requires_grad_(True)
+        interpolated_logits = self.critic(interpolated, conditioned_images)
+        
         gradients = torch.autograd.grad(outputs=interpolated_logits, inputs=interpolated,
                                         grad_outputs=torch.ones_like(interpolated_logits), create_graph=True, retain_graph=True)[0]
         gradients = gradients.view(len(gradients), -1)
         gradients_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
         loss_C += self.lambda_gp * gradients_penalty
-
+        
         # Compute the R1 regularization loss
         r1_reg = gradients.pow(2).sum(1).mean()
         loss_C += self.lambda_r1 * r1_reg
@@ -122,53 +116,65 @@ class CWGAN(pl.LightningModule):
         loss_C.backward()
         optimizer_C.step()
         self.critic_losses.append(loss_C.item())
-
+    
     def training_step(self, batch, batch_idx):
-        optimizer_C, optimizer_G = self.optimizers()
-        real_images, conditioned_images = batch
-
-        self.critic_step(real_images, conditioned_images, optimizer_C)
-        self.generator_step(real_images, conditioned_images, optimizer_G)
-
+        real, condition = batch
+        optimizer_C, optimizer_G = self.optimizers()  # Access optimizers
+    
+        # Update the critic
+        self.critic_step(real, condition, optimizer_C)
+    
+        # Update the generator
+        self.generator_step(real, condition, optimizer_G)
+    
+        # Logging and saving models
+        gen_mean = sum(self.generator_losses[-self.display_step:]) / self.display_step
+        crit_mean = sum(self.critic_losses[-self.display_step:]) / self.display_step
         if self.current_epoch % self.display_step == 0 and batch_idx == 0:
-            with torch.no_grad():
-                fake_images = self.generator(conditioned_images)
-                display_progress(conditioned_images[0], real_images[0], fake_images[0], self.current_epoch)
-
-    def on_epoch_end(self):
-        gc.collect()
-        torch.cuda.empty_cache()
-
-    def forward(self, x):
-        return self.generator(x)
+            fake = self.generator(condition).detach()
+            logger.info(f"Epoch {self.current_epoch}: Generator loss: {gen_mean}, Critic loss: {crit_mean}")
+            display_progress(condition[0], real[0], fake[0], self.current_epoch)
     
 
 import os
 class ModelTrainer:
-    def __init__(self, config : ModelTrainerConfig):
+    def __init__(self, config):
         self.config = config
         
     def load_datasets(self):
         self.train_dataset = torch.load(self.config.train_data_path)
         self.test_dataset = torch.load(self.config.test_data_path)
+    
+        # Ensure these are actually datasets, not dataloaders
+        if isinstance(self.train_dataset, DataLoader):
+            self.train_dataset = self.train_dataset.dataset
+        if isinstance(self.test_dataset, DataLoader):
+            self.test_dataset = self.test_dataset.dataset
         
     def create_dataloaders(self):
         self.train_dataloader = DataLoader(
-            self.train_dataset, batch_size=self.config.BATCH_SIZE, shuffle=True
+            self.train_dataset, 
+            batch_size=self.config.BATCH_SIZE, 
+            shuffle=True,
         )
         self.test_dataloader = DataLoader(
-            self.test_dataset, batch_size=self.config.BATCH_SIZE, shuffle=False
+            self.test_dataset, 
+            batch_size=self.config.BATCH_SIZE, 
+            shuffle=False,
         )
+    
+      
+        
+        
         
     def initialize_model(self):
         self.model = CWGAN(
-            in_channels=self.config.INPUT_CHANNELS,
-            out_channels=self.config.OUTPUT_CHANNELS,
+            in_channels=1,
+            out_channels=2,
             learning_rate=self.config.LEARNING_RATE,
             lambda_recon=self.config.LAMBDA_RECON,
             display_step=self.config.DISPLAY_STEP,
-            lambda_gp=10,  # Default value, you can make it configurable
-            lambda_r1=10   # Default value, you can make it configurable
+            
         )
         
     def train_model(self):

src/imagecolorization/pipeline/stage_04_model_trainer.py

@@ -14,4 +14,4 @@ class ModelTrainerPipeline:
         model_trainer.create_dataloaders()
         model_trainer.initialize_model()
         model_trainer.train_model()
-        model_trainer.save_model()
+        model_trainer.save_model()