bug fixes

CatVTON_model.py

@@ -0,0 +1,138 @@
+import inspect
+import os
+from typing import Union
+
+import PIL
+import numpy as np
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from utils import (check_inputs_maskfree, get_time_embedding, numpy_to_pil, prepare_image, compute_vae_encodings)
+from ddpm import DDPMSampler
+from tqdm import tqdm
+
+
+class CatVTONPix2PixPipeline:
+    def __init__(
+        self, 
+        weight_dtype=torch.float32,
+        device='cuda',
+        compile=False,
+        skip_safety_check=True,
+        use_tf32=True,
+        models={},
+    ):
+        self.device = device
+        self.weight_dtype = weight_dtype
+        self.skip_safety_check = skip_safety_check
+        self.models = models
+
+        self.generator = torch.Generator(device=device)
+        self.noise_scheduler = DDPMSampler(generator=self.generator)
+        # self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device, dtype=weight_dtype)
+        self.encoder= models.get('encoder', None)
+        self.decoder= models.get('decoder', None)
+ 
+        self.unet=models.get('diffusion', None)  
+        # # Enable TF32 for faster training on Ampere GPUs (A100 and RTX 30 series).
+        if use_tf32:
+            torch.set_float32_matmul_precision("high")
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+    @torch.no_grad()
+    def __call__(
+        self, 
+        image: Union[PIL.Image.Image, torch.Tensor],
+        condition_image: Union[PIL.Image.Image, torch.Tensor],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 2.5,
+        height: int = 1024,
+        width: int = 768,
+        generator=None,
+        eta=1.0,
+        **kwargs
+    ):
+        concat_dim = -1  # FIXME: y axis concat
+        # Prepare inputs to Tensor
+        image, condition_image = check_inputs_maskfree(image, condition_image, width, height)
+        
+        image = prepare_image(image).to(self.device, dtype=self.weight_dtype)
+        
+        condition_image = prepare_image(condition_image).to(self.device, dtype=self.weight_dtype)
+        
+        # Encode the image
+        image_latent = compute_vae_encodings(image, self.encoder)
+        condition_latent = compute_vae_encodings(condition_image, self.encoder)
+        
+        del image, condition_image
+        # Concatenate latents
+        # Concatenate latents
+        condition_latent_concat = torch.cat([image_latent, condition_latent], dim=concat_dim)
+        # Prepare noise
+        latents = randn_tensor(
+            condition_latent_concat.shape,
+            generator=generator,
+            device=condition_latent_concat.device,
+            dtype=self.weight_dtype,
+        )
+        # Prepare timesteps
+        self.noise_scheduler.set_inference_timesteps(num_inference_steps)
+        timesteps = self.noise_scheduler.timesteps
+        # latents = latents * self.noise_scheduler.init_noise_sigma
+        latents = self.noise_scheduler.add_noise(latents, timesteps[0])
+        
+        # Classifier-Free Guidance
+        if do_classifier_free_guidance := (guidance_scale > 1.0):
+            condition_latent_concat = torch.cat(
+                [
+                    torch.cat([image_latent, torch.zeros_like(condition_latent)], dim=concat_dim),
+                    condition_latent_concat,
+                ]
+            )
+
+        num_warmup_steps = 0  # For simple DDPM, no warmup needed
+        with tqdm(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                
+                latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents)
+
+                # prepare the input for the inpainting model
+                
+                p2p_latent_model_input = torch.cat([latent_model_input, condition_latent_concat], dim=1)
+                # predict the noise residual
+                
+                timestep = t.repeat(p2p_latent_model_input.shape[0])
+                time_embedding = get_time_embedding(timestep).to(self.device, dtype=self.weight_dtype)
+
+                noise_pred = self.unet(
+                    p2p_latent_model_input,
+                    time_embedding
+                )
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.noise_scheduler.step(
+                    t, latents, noise_pred
+                )
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps
+                ):
+                    progress_bar.update()
+
+        # Decode the final latents
+        latents = latents.split(latents.shape[concat_dim] // 2, dim=concat_dim)[0]
+        # latents = 1 / self.vae.config.scaling_factor * latents
+        # image = self.vae.decode(latents.to(self.device, dtype=self.weight_dtype)).sample
+        image = self.decoder(latents.to(self.device, dtype=self.weight_dtype))
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        image = numpy_to_pil(image)
+        
+        return image

app.py

@@ -8,146 +8,32 @@ from typing import Optional
 # Import your custom modules
 from load_model import preload_models_from_standard_weights
 from utils import to_pil_image
+from CatVTON_model import CatVTONPix2PixPipeline
 
-import inspect
 import os
-from typing import Union
-
-import PIL
-import numpy as np
 import torch
-import tqdm
-from diffusers.utils.torch_utils import randn_tensor
-
-from utils import (check_inputs_maskfree, get_time_embedding, numpy_to_pil, prepare_image, compute_vae_encodings)
-from ddpm import DDPMSampler
-
-
-class CatVTONPix2PixPipeline:
-    def __init__(
-        self, 
-        weight_dtype=torch.float32,
-        device='cuda',
-        compile=False,
-        skip_safety_check=True,
-        use_tf32=True,
-        models={},
-    ):
-        self.device = device
-        self.weight_dtype = weight_dtype
-        self.skip_safety_check = skip_safety_check
-        self.models = models
+import urllib.request
 
-        self.generator = torch.Generator(device=device)
-        self.noise_scheduler = DDPMSampler(generator=self.generator)
-        self.encoder= models.get('encoder', None)
-        self.decoder= models.get('decoder', None)
-        self.unet=models.get('diffusion', None)  
-        
-        # Enable TF32 for faster training on Ampere GPUs
-        if use_tf32:
-            torch.set_float32_matmul_precision("high")
-            torch.backends.cuda.matmul.allow_tf32 = True
-
-    @torch.no_grad()
-    def __call__(
-        self, 
-        image: Union[PIL.Image.Image, torch.Tensor],
-        condition_image: Union[PIL.Image.Image, torch.Tensor],
-        num_inference_steps: int = 50,
-        guidance_scale: float = 2.5,
-        height: int = 1024,
-        width: int = 768,
-        generator=None,
-        eta=1.0,
-        **kwargs
-    ):
-        concat_dim = -1  # FIXME: y axis concat
-        # Prepare inputs to Tensor
-        image, condition_image = check_inputs_maskfree(image, condition_image, width, height)
-        
-        # Ensure consistent dtype for all tensors
-        image = prepare_image(image).to(self.device, dtype=self.weight_dtype)
-        condition_image = prepare_image(condition_image).to(self.device, dtype=self.weight_dtype)
-        
-        # Encode the image
-        image_latent = compute_vae_encodings(image, self.encoder)
-        condition_latent = compute_vae_encodings(condition_image, self.encoder)
-        
-        del image, condition_image
-        
-        # Concatenate latents
-        condition_latent_concat = torch.cat([image_latent, condition_latent], dim=concat_dim)
-        
-        # Prepare noise
-        latents = randn_tensor(
-            condition_latent_concat.shape,
-            generator=generator,
-            device=condition_latent_concat.device,
-            dtype=self.weight_dtype,
-        )
-        
-        # Prepare timesteps
-        self.noise_scheduler.set_inference_timesteps(num_inference_steps)
-        timesteps = self.noise_scheduler.timesteps
-        latents = self.noise_scheduler.add_noise(latents, timesteps[0])
-        
-        # Classifier-Free Guidance
-        if do_classifier_free_guidance := (guidance_scale > 1.0):
-            condition_latent_concat = torch.cat(
-                [
-                    torch.cat([image_latent, torch.zeros_like(condition_latent)], dim=concat_dim),
-                    condition_latent_concat,
-                ]
-            )
+def load_models():
+    try:
+        print("🚀 Starting model loading process...")
 
-        num_warmup_steps = 0  # For simple DDPM, no warmup needed
-        with tqdm.tqdm(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents)
+        # Check and download model files if missing
+        ckpt_path = "instruct-pix2pix-00-22000.ckpt"
+        finetune_path = "maskfree_finetuned_weights.safetensors"
 
-                # prepare the input for the inpainting model
-                p2p_latent_model_input = torch.cat([latent_model_input, condition_latent_concat], dim=1)
-                
-                # predict the noise residual
-                timestep = t.repeat(p2p_latent_model_input.shape[0])
-                time_embedding = get_time_embedding(timestep).to(self.device, dtype=self.weight_dtype)
+        if not os.path.exists(ckpt_path):
+            print(f"⬇️ Downloading {ckpt_path}...")
+            url = "https://huggingface.co/timbrooks/instruct-pix2pix/resolve/main/instruct-pix2pix-00-22000.ckpt"
+            urllib.request.urlretrieve(url, ckpt_path)
+            print("✅ Download complete.")
+        else:
+            print("✅ Checkpoint already exists.")
 
-                noise_pred = self.unet(
-                    p2p_latent_model_input,
-                    time_embedding
-                )
-                
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (
-                        noise_pred_text - noise_pred_uncond
-                    )
-                
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.noise_scheduler.step(
-                    t, latents, noise_pred
-                )
-                
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps):
-                    progress_bar.update()
+        if not os.path.exists(finetune_path):
+            print(f"❌ Finetune weights file not found: {finetune_path}")
+            return None, None
 
-        # Decode the final latents
-        latents = latents.split(latents.shape[concat_dim] // 2, dim=concat_dim)[0]
-        image = self.decoder(latents.to(self.device, dtype=self.weight_dtype))
-        image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        image = numpy_to_pil(image)
-        
-        return image
-    
-def load_models():
-    try:
-        print("🚀 Starting model loading process...")
-        
         # Check CUDA availability
         cuda_available = torch.cuda.is_available()
         print(f"CUDA available: {cuda_available}")
@@ -155,21 +41,9 @@ def load_models():
             print(f"CUDA device: {torch.cuda.get_device_name()}")
             free_memory = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)
             print(f"Available CUDA memory: {free_memory / 1e9:.2f} GB")
-        
+
         device = "cuda" if cuda_available else "cpu"
-        
-        # Check if model files exist
-        ckpt_path = "instruct-pix2pix-00-22000.ckpt"
-        finetune_path = "maskfree_finetuned_weights.safetensors"
-        
-        if not os.path.exists(ckpt_path):
-            print(f"❌ Checkpoint file not found: {ckpt_path}")
-            return None, None
-            
-        if not os.path.exists(finetune_path):
-            print(f"❌ Finetune weights file not found: {finetune_path}")
-            return None, None
-        
+
         print("📦 Loading models from weights...")
         
         models = preload_models_from_standard_weights(
@@ -177,16 +51,14 @@ def load_models():
             device=device, 
             finetune_weights_path=finetune_path
         )
-        
+
         if not models:
             print("❌ Failed to load models")
             return None, None
-        
-        # Convert all models to consistent dtype to avoid mixed precision issues
-        weight_dtype = torch.float32  # Use float32 to avoid dtype mismatch
+
+        weight_dtype = torch.float32
         print(f"Converting models to {weight_dtype}...")
-        
-        # Ensure all models use the same dtype
+
         for model_name, model in models.items():
             if model is not None:
                 try:
@@ -195,19 +67,19 @@ def load_models():
                     print(f"✅ {model_name} converted to {weight_dtype}")
                 except Exception as e:
                     print(f"⚠️ Could not convert {model_name} to {weight_dtype}: {e}")
-        
+
         print("🔧 Initializing pipeline...")
-        
+
         pipeline = CatVTONPix2PixPipeline(
             weight_dtype=weight_dtype,
             device=device,
             skip_safety_check=True,
             models=models,
         )
-        
+
         print("✅ Models and pipeline loaded successfully!")
         return models, pipeline
-        
+
     except Exception as e:
         print(f"❌ Error in load_models: {e}")
         import traceback

mask_free_inference.ipynb

@@ -95,197 +95,39 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "78e3d8b9",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/mahesh/miniconda3/envs/harsh/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    }
-   ],
-   "source": [
-    "import inspect\n",
-    "import os\n",
-    "from typing import Union\n",
-    "\n",
-    "import PIL\n",
-    "import numpy as np\n",
-    "import torch\n",
-    "import tqdm\n",
-    "from diffusers.utils.torch_utils import randn_tensor\n",
-    "\n",
-    "from utils import (check_inputs_maskfree, get_time_embedding, numpy_to_pil, prepare_image, compute_vae_encodings)\n",
-    "from ddpm import DDPMSampler\n",
-    "\n",
-    "class CatVTONPix2PixPipeline:\n",
-    "    def __init__(\n",
-    "        self, \n",
-    "        weight_dtype=torch.float32,\n",
-    "        device='cuda',\n",
-    "        compile=False,\n",
-    "        skip_safety_check=True,\n",
-    "        use_tf32=True,\n",
-    "        models={},\n",
-    "    ):\n",
-    "        self.device = device\n",
-    "        self.weight_dtype = weight_dtype\n",
-    "        self.skip_safety_check = skip_safety_check\n",
-    "        self.models = models\n",
-    "\n",
-    "        self.generator = torch.Generator(device=device)\n",
-    "        self.noise_scheduler = DDPMSampler(generator=self.generator)\n",
-    "        # self.vae = AutoencoderKL.from_pretrained(\"stabilityai/sd-vae-ft-mse\").to(device, dtype=weight_dtype)\n",
-    "        self.encoder= models.get('encoder', None)\n",
-    "        self.decoder= models.get('decoder', None)\n",
-    " \n",
-    "        self.unet=models.get('diffusion', None)  \n",
-    "        # # Enable TF32 for faster training on Ampere GPUs (A100 and RTX 30 series).\n",
-    "        if use_tf32:\n",
-    "            torch.set_float32_matmul_precision(\"high\")\n",
-    "            torch.backends.cuda.matmul.allow_tf32 = True\n",
-    "\n",
-    "    @torch.no_grad()\n",
-    "    def __call__(\n",
-    "        self, \n",
-    "        image: Union[PIL.Image.Image, torch.Tensor],\n",
-    "        condition_image: Union[PIL.Image.Image, torch.Tensor],\n",
-    "        num_inference_steps: int = 50,\n",
-    "        guidance_scale: float = 2.5,\n",
-    "        height: int = 1024,\n",
-    "        width: int = 768,\n",
-    "        generator=None,\n",
-    "        eta=1.0,\n",
-    "        **kwargs\n",
-    "    ):\n",
-    "        concat_dim = -1  # FIXME: y axis concat\n",
-    "        # Prepare inputs to Tensor\n",
-    "        image, condition_image = check_inputs_maskfree(image, condition_image, width, height)\n",
-    "        \n",
-    "        image = prepare_image(image).to(self.device, dtype=self.weight_dtype)\n",
-    "        \n",
-    "        condition_image = prepare_image(condition_image).to(self.device, dtype=self.weight_dtype)\n",
-    "        \n",
-    "        # Encode the image\n",
-    "        image_latent = compute_vae_encodings(image, self.encoder)\n",
-    "        condition_latent = compute_vae_encodings(condition_image, self.encoder)\n",
-    "        \n",
-    "        del image, condition_image\n",
-    "        # Concatenate latents\n",
-    "        # Concatenate latents\n",
-    "        condition_latent_concat = torch.cat([image_latent, condition_latent], dim=concat_dim)\n",
-    "        # Prepare noise\n",
-    "        latents = randn_tensor(\n",
-    "            condition_latent_concat.shape,\n",
-    "            generator=generator,\n",
-    "            device=condition_latent_concat.device,\n",
-    "            dtype=self.weight_dtype,\n",
-    "        )\n",
-    "        # Prepare timesteps\n",
-    "        self.noise_scheduler.set_inference_timesteps(num_inference_steps)\n",
-    "        timesteps = self.noise_scheduler.timesteps\n",
-    "        # latents = latents * self.noise_scheduler.init_noise_sigma\n",
-    "        latents = self.noise_scheduler.add_noise(latents, timesteps[0])\n",
-    "        \n",
-    "        # Classifier-Free Guidance\n",
-    "        if do_classifier_free_guidance := (guidance_scale > 1.0):\n",
-    "            condition_latent_concat = torch.cat(\n",
-    "                [\n",
-    "                    torch.cat([image_latent, torch.zeros_like(condition_latent)], dim=concat_dim),\n",
-    "                    condition_latent_concat,\n",
-    "                ]\n",
-    "            )\n",
-    "\n",
-    "        num_warmup_steps = 0  # For simple DDPM, no warmup needed\n",
-    "        with tqdm(total=num_inference_steps) as progress_bar:\n",
-    "            for i, t in enumerate(timesteps):\n",
-    "                # expand the latents if we are doing classifier free guidance\n",
-    "                \n",
-    "                latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents)\n",
-    "\n",
-    "                # prepare the input for the inpainting model\n",
-    "                \n",
-    "                p2p_latent_model_input = torch.cat([latent_model_input, condition_latent_concat], dim=1)\n",
-    "                # predict the noise residual\n",
-    "                \n",
-    "                timestep = t.repeat(p2p_latent_model_input.shape[0])\n",
-    "                time_embedding = get_time_embedding(timestep).to(self.device, dtype=self.weight_dtype)\n",
-    "\n",
-    "                noise_pred = self.unet(\n",
-    "                    p2p_latent_model_input,\n",
-    "                    time_embedding\n",
-    "                )\n",
-    "                # perform guidance\n",
-    "                if do_classifier_free_guidance:\n",
-    "                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
-    "                    noise_pred = noise_pred_uncond + guidance_scale * (\n",
-    "                        noise_pred_text - noise_pred_uncond\n",
-    "                    )\n",
-    "                # compute the previous noisy sample x_t -> x_t-1\n",
-    "                latents = self.noise_scheduler.step(\n",
-    "                    t, latents, noise_pred\n",
-    "                )\n",
-    "                # call the callback, if provided\n",
-    "                if i == len(timesteps) - 1 or (\n",
-    "                    (i + 1) > num_warmup_steps\n",
-    "                ):\n",
-    "                    progress_bar.update()\n",
-    "\n",
-    "        # Decode the final latents\n",
-    "        latents = latents.split(latents.shape[concat_dim] // 2, dim=concat_dim)[0]\n",
-    "        # latents = 1 / self.vae.config.scaling_factor * latents\n",
-    "        # image = self.vae.decode(latents.to(self.device, dtype=self.weight_dtype)).sample\n",
-    "        image = self.decoder(latents.to(self.device, dtype=self.weight_dtype))\n",
-    "        image = (image / 2 + 0.5).clamp(0, 1)\n",
-    "        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16\n",
-    "        image = image.cpu().permute(0, 2, 3, 1).float().numpy()\n",
-    "        image = numpy_to_pil(image)\n",
-    "        \n",
-    "        return image\n"
-   ]
+   "outputs": [],
+   "source": []
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "5627b2d2",
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Dataset vitonhd loaded, total 20 pairs.\n"
+      "/home/mahesh/miniconda3/envs/harsh/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
      ]
     },
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 50/50 [00:07<00:00,  7.12it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.31it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  7.09it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.98it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  7.01it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  7.13it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.28it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  7.13it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.17it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.97it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.17it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.38it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.20it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.92it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.71it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.25it/s]\n",
-      "100%|██████████| 50/50 [00:06<00:00,  7.49it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.87it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.89it/s]\n",
-      "100%|██████████| 50/50 [00:07<00:00,  6.92it/s]\n",
-      "100%|██████████| 20/20 [02:26<00:00,  7.35s/it]\n"
+     "ename": "AssertionError",
+     "evalue": "File ./sample_dataset/samples_pairs.txt does not exist.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[3], line 97\u001b[0m\n\u001b[1;32m     94\u001b[0m             result\u001b[38;5;241m.\u001b[39msave(output_path)\n\u001b[1;32m     96\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 97\u001b[0m     \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/harsh/lib/python3.10/site-packages/torch/utils/_contextlib.py:116\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    114\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    115\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 116\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\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[0;32mIn[3], line 48\u001b[0m, in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[38;5;66;03m# Dataset\u001b[39;00m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m args\u001b[38;5;241m.\u001b[39mdataset_name \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvitonhd\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 48\u001b[0m     dataset \u001b[38;5;241m=\u001b[39m \u001b[43mVITONHDTestDataset\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid dataset name \u001b[39m\u001b[38;5;132;01m{\u001b[39;00margs\u001b[38;5;241m.\u001b[39mdataset\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "File \u001b[0;32m~/harsh/stable-diffusion/VITON_Dataset.py:14\u001b[0m, in \u001b[0;36mInferenceDataset.__init__\u001b[0;34m(self, args)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvae_processor \u001b[38;5;241m=\u001b[39m VaeImageProcessor(vae_scale_factor\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m8\u001b[39m) \n\u001b[1;32m     13\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmask_processor \u001b[38;5;241m=\u001b[39m VaeImageProcessor(vae_scale_factor\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m8\u001b[39m, do_normalize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, do_binarize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, do_convert_grayscale\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \n\u001b[0;32m---> 14\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdata \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/harsh/stable-diffusion/VITON_Dataset.py:36\u001b[0m, in \u001b[0;36mVITONHDTestDataset.load_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mload_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m     35\u001b[0m     name\u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mis_train \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msamples\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m---> 36\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mexists(pair_txt\u001b[38;5;241m:=\u001b[39mos\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mdata_root_path, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_pairs.txt\u001b[39m\u001b[38;5;124m'\u001b[39m)), \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFile \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpair_txt\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m does not exist.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     37\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(pair_txt, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[1;32m     38\u001b[0m         lines \u001b[38;5;241m=\u001b[39m f\u001b[38;5;241m.\u001b[39mreadlines()\n",
+      "\u001b[0;31mAssertionError\u001b[0m: File ./sample_dataset/samples_pairs.txt does not exist."
      ]
     }
    ],
@@ -297,6 +139,7 @@
     "from VITON_Dataset import VITONHDTestDataset\n",
     "from tqdm import tqdm\n",
     "from PIL import Image\n",
+    "from CatVTON_model import CatVTONPix2PixPipeline\n",
     "\n",
     "from utils import to_pil_image\n",
     "\n",