bug fixes

VITON_Dataset.py

@@ -32,7 +32,7 @@ class InferenceDataset(Dataset):
 
 class VITONHDTestDataset(InferenceDataset):
     def load_data(self):
-        name= "train" if self.args.is_train else "test"
+        name= "train" if self.args.is_train else "samples"
         assert os.path.exists(pair_txt:=os.path.join(self.args.data_root_path, f'{name}_pairs.txt')), f"File {pair_txt} does not exist."
         with open(pair_txt, 'r') as f:
             lines = f.readlines()

interface.py

@@ -4,7 +4,7 @@ from PIL import Image
 from transformers import CLIPTokenizer
 
 # Import your existing model and pipeline modules
-import model
+import load_model
 import pipeline
 
 # Device Configuration
@@ -24,7 +24,7 @@ print(f"Using device: {DEVICE}")
 # Load tokenizer and models
 tokenizer = CLIPTokenizer("vocab.json", merges_file="merges.txt")
 model_file = "inkpunk-diffusion-v1.ckpt"
-models = model.preload_models_from_standard_weights(model_file, DEVICE)
+models = load_model.preload_models_from_standard_weights(model_file, DEVICE)
 # models=None
 
 def generate_image(

pipeline.py

@@ -6,7 +6,7 @@ import numpy as np
 from tqdm import tqdm
 from ddpm import DDPMSampler
 from PIL import Image
-import model
+import load_model
 from utils import check_inputs, prepare_image, prepare_mask_image
 
 WIDTH = 512
@@ -293,7 +293,7 @@ if __name__ == "__main__":
     mask = Image.open("agnostic_mask.png").convert("L")
 
     # Load models
-    models=model.preload_models_from_standard_weights("sd-v1-5-inpainting.ckpt", device="cuda")
+    models=load_model.preload_models_from_standard_weights("sd-v1-5-inpainting.ckpt", device="cuda")
 
     # Generate image
     generated_image = generate(

test.ipynb

@@ -2,67 +2,52 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "id": "867520bc",
+   "execution_count": null,
+   "id": "6387c9e1",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Cloning into 'CatVTON'...\n",
-      "remote: Enumerating objects: 1348, done.\u001b[K\n",
-      "remote: Counting objects: 100% (62/62), done.\u001b[K\n",
-      "remote: Compressing objects: 100% (29/29), done.\u001b[K\n",
-      "remote: Total 1348 (delta 51), reused 33 (delta 33), pack-reused 1286 (from 3)\u001b[K\n",
-      "Receiving objects: 100% (1348/1348), 16.74 MiB | 42.65 MiB/s, done.\n",
-      "Resolving deltas: 100% (449/449), done.\n"
-     ]
-    }
-   ],
-   "source": [
-    "!git clone https://github.com/Zheng-Chong/CatVTON.git"
-   ]
+   "outputs": [],
+   "source": []
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "id": "3d2f98af",
+   "execution_count": null,
+   "id": "ca9233f0",
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[0m\u001b[01;34mtest\u001b[0m/  test_pairs.txt  \u001b[01;34mtrain\u001b[0m/  train_pairs.txt\n"
-     ]
+     "data": {
+      "text/plain": [
+       "'/kaggle/working'"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
-   "source": [
-    "ls /kaggle/input/viton-hd-dataset"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
-   "id": "ba750da0",
+   "execution_count": 17,
+   "id": "3d2f98af",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "/kaggle/working/stable-diffusion/CatVTON\n"
+      "\u001b[0m\u001b[01;34mtest\u001b[0m/  \u001b[01;32mtest_pairs.txt\u001b[0m*  \u001b[01;34mtrain\u001b[0m/  \u001b[01;32mtrain_pairs.txt\u001b[0m*\n"
      ]
     }
    ],
    "source": [
-    "cd CatVTON/"
+    "ls /kaggle/input/viton-hd-dataset"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 18,
    "id": "dc0f36f4",
    "metadata": {},
    "outputs": [
@@ -71,12 +56,12 @@
      "output_type": "stream",
      "text": [
       "Cloning into 'stable-diffusion'...\n",
-      "remote: Enumerating objects: 56, done.\u001b[K\n",
-      "remote: Counting objects: 100% (56/56), done.\u001b[K\n",
-      "remote: Compressing objects: 100% (44/44), done.\u001b[K\n",
-      "remote: Total 56 (delta 17), reused 50 (delta 12), pack-reused 0 (from 0)\u001b[K\n",
-      "Receiving objects: 100% (56/56), 4.68 MiB | 36.31 MiB/s, done.\n",
-      "Resolving deltas: 100% (17/17), done.\n"
+      "remote: Enumerating objects: 150, done.\u001b[K\n",
+      "remote: Counting objects: 100% (150/150), done.\u001b[K\n",
+      "remote: Compressing objects: 100% (124/124), done.\u001b[K\n",
+      "remote: Total 150 (delta 36), reused 139 (delta 26), pack-reused 0 (from 0)\u001b[K\n",
+      "Receiving objects: 100% (150/150), 9.11 MiB | 20.74 MiB/s, done.\n",
+      "Resolving deltas: 100% (36/36), done.\n"
      ]
     }
    ],
@@ -86,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 19,
    "id": "a0bf01ab",
    "metadata": {},
    "outputs": [
@@ -94,7 +79,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "/kaggle/working/stable-diffusion/CatVTON/stable-diffusion\n"
+      "/kaggle/working/stable-diffusion\n"
      ]
     }
    ],
@@ -104,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 20,
    "id": "1401cd56",
    "metadata": {},
    "outputs": [
@@ -112,25 +97,25 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "--2025-06-11 10:33:00--  https://huggingface.co/sd-legacy/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt\n",
-      "Resolving huggingface.co (huggingface.co)... 3.163.189.114, 3.163.189.74, 3.163.189.90, ...\n",
-      "Connecting to huggingface.co (huggingface.co)|3.163.189.114|:443... connected.\n",
+      "--2025-06-13 07:07:34--  https://huggingface.co/sd-legacy/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt\n",
+      "Resolving huggingface.co (huggingface.co)... 18.67.93.22, 18.67.93.63, 18.67.93.58, ...\n",
+      "Connecting to huggingface.co (huggingface.co)|18.67.93.22|:443... connected.\n",
       "HTTP request sent, awaiting response... 307 Temporary Redirect\n",
       "Location: /stable-diffusion-v1-5/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt [following]\n",
-      "--2025-06-11 10:33:01--  https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt\n",
+      "--2025-06-13 07:07:34--  https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt\n",
       "Reusing existing connection to huggingface.co:443.\n",
       "HTTP request sent, awaiting response... 302 Found\n",
-      "Location: https://cdn-lfs.hf.co/repos/f6/56/f656f0fa3b8a40ac76d297fa2a4b00f981e8eb1261963460764e7dd3b35ec97f/c6bbc15e3224e6973459ba78de4998b80b50112b0ae5b5c67113d56b4e366b19?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27sd-v1-5-inpainting.ckpt%3B+filename%3D%22sd-v1-5-inpainting.ckpt%22%3B&Expires=1749640621&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc0OTY0MDYyMX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5oZi5jby9yZXBvcy9mNi81Ni9mNjU2ZjBmYTNiOGE0MGFjNzZkMjk3ZmEyYTRiMDBmOTgxZThlYjEyNjE5NjM0NjA3NjRlN2RkM2IzNWVjOTdmL2M2YmJjMTVlMzIyNGU2OTczNDU5YmE3OGRlNDk5OGI4MGI1MDExMmIwYWU1YjVjNjcxMTNkNTZiNGUzNjZiMTk%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=sdomKXQPt3COUrAxFqHQCR83b0Sgw0kHwStFv%7EqgSrCfwOddw9sNRX3qela0jgons998TT3Oqk0TA0c-PTLyPpAO-iqM9aGvLsRLixtxMNgdvDdWdk87Ywdgvg24T6GkVaL9I8ErFlF918m%7EYMtHICZ8hcoq1GST-DdDigp4vA-w9lHnRfOGteBzViPKyqgQaYiYRd10FVmSYYpFUJrZ%7ECFAGO5MwVA-OTlMVLOYKKPs0s3duoP4KIz9-SUoUIXbgUmiuExLqdVulk-tJRCSAk-u7WvbUhPUsraiP1YGa-QvUYoygX5xlluuFIt%7EG54t5TrCzIWP0tu0ZGaqr3%7E%7EEA__&Key-Pair-Id=K3RPWS32NSSJCE [following]\n",
-      "--2025-06-11 10:33:01--  https://cdn-lfs.hf.co/repos/f6/56/f656f0fa3b8a40ac76d297fa2a4b00f981e8eb1261963460764e7dd3b35ec97f/c6bbc15e3224e6973459ba78de4998b80b50112b0ae5b5c67113d56b4e366b19?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27sd-v1-5-inpainting.ckpt%3B+filename%3D%22sd-v1-5-inpainting.ckpt%22%3B&Expires=1749640621&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc0OTY0MDYyMX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5oZi5jby9yZXBvcy9mNi81Ni9mNjU2ZjBmYTNiOGE0MGFjNzZkMjk3ZmEyYTRiMDBmOTgxZThlYjEyNjE5NjM0NjA3NjRlN2RkM2IzNWVjOTdmL2M2YmJjMTVlMzIyNGU2OTczNDU5YmE3OGRlNDk5OGI4MGI1MDExMmIwYWU1YjVjNjcxMTNkNTZiNGUzNjZiMTk%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=sdomKXQPt3COUrAxFqHQCR83b0Sgw0kHwStFv%7EqgSrCfwOddw9sNRX3qela0jgons998TT3Oqk0TA0c-PTLyPpAO-iqM9aGvLsRLixtxMNgdvDdWdk87Ywdgvg24T6GkVaL9I8ErFlF918m%7EYMtHICZ8hcoq1GST-DdDigp4vA-w9lHnRfOGteBzViPKyqgQaYiYRd10FVmSYYpFUJrZ%7ECFAGO5MwVA-OTlMVLOYKKPs0s3duoP4KIz9-SUoUIXbgUmiuExLqdVulk-tJRCSAk-u7WvbUhPUsraiP1YGa-QvUYoygX5xlluuFIt%7EG54t5TrCzIWP0tu0ZGaqr3%7E%7EEA__&Key-Pair-Id=K3RPWS32NSSJCE\n",
-      "Resolving cdn-lfs.hf.co (cdn-lfs.hf.co)... 18.172.170.108, 18.172.170.21, 18.172.170.5, ...\n",
-      "Connecting to cdn-lfs.hf.co (cdn-lfs.hf.co)|18.172.170.108|:443... connected.\n",
+      "Location: https://cdn-lfs.hf.co/repos/f6/56/f656f0fa3b8a40ac76d297fa2a4b00f981e8eb1261963460764e7dd3b35ec97f/c6bbc15e3224e6973459ba78de4998b80b50112b0ae5b5c67113d56b4e366b19?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27sd-v1-5-inpainting.ckpt%3B+filename%3D%22sd-v1-5-inpainting.ckpt%22%3B&Expires=1749802055&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc0OTgwMjA1NX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5oZi5jby9yZXBvcy9mNi81Ni9mNjU2ZjBmYTNiOGE0MGFjNzZkMjk3ZmEyYTRiMDBmOTgxZThlYjEyNjE5NjM0NjA3NjRlN2RkM2IzNWVjOTdmL2M2YmJjMTVlMzIyNGU2OTczNDU5YmE3OGRlNDk5OGI4MGI1MDExMmIwYWU1YjVjNjcxMTNkNTZiNGUzNjZiMTk%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=m4Xzc4SaPX28SXT9wK8qPXBWIr7uKmVt6iq2D3qMALrJWCfm1a4FHoshqkXLDrchchDIkAImr7l-yDlAv9x15JsX09FidLsSVU8UXS4a%7Em69hgWMTgloTObR3HlTwY9EQ7t%7ErneASRUS5r%7E2szyfyrlN-n4-U9QWCmyOikaumCc0PbAHE6lRNcy7FSCTxQGM48h%7EQBZ37iQArWW2JC%7E-apwm1knzGt422ywPlQws2qREoUeCPoXFWKl-iX1%7EqDimjSepdm2ZGt-COfekmJddQWXuCQAj7uY5YKcE3qEt7IBcaj96MNbF8b2qxTNbLrzgXioIzl0SIw8Ws-YUOu5I3A__&Key-Pair-Id=K3RPWS32NSSJCE [following]\n",
+      "--2025-06-13 07:07:35--  https://cdn-lfs.hf.co/repos/f6/56/f656f0fa3b8a40ac76d297fa2a4b00f981e8eb1261963460764e7dd3b35ec97f/c6bbc15e3224e6973459ba78de4998b80b50112b0ae5b5c67113d56b4e366b19?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27sd-v1-5-inpainting.ckpt%3B+filename%3D%22sd-v1-5-inpainting.ckpt%22%3B&Expires=1749802055&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc0OTgwMjA1NX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5oZi5jby9yZXBvcy9mNi81Ni9mNjU2ZjBmYTNiOGE0MGFjNzZkMjk3ZmEyYTRiMDBmOTgxZThlYjEyNjE5NjM0NjA3NjRlN2RkM2IzNWVjOTdmL2M2YmJjMTVlMzIyNGU2OTczNDU5YmE3OGRlNDk5OGI4MGI1MDExMmIwYWU1YjVjNjcxMTNkNTZiNGUzNjZiMTk%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=m4Xzc4SaPX28SXT9wK8qPXBWIr7uKmVt6iq2D3qMALrJWCfm1a4FHoshqkXLDrchchDIkAImr7l-yDlAv9x15JsX09FidLsSVU8UXS4a%7Em69hgWMTgloTObR3HlTwY9EQ7t%7ErneASRUS5r%7E2szyfyrlN-n4-U9QWCmyOikaumCc0PbAHE6lRNcy7FSCTxQGM48h%7EQBZ37iQArWW2JC%7E-apwm1knzGt422ywPlQws2qREoUeCPoXFWKl-iX1%7EqDimjSepdm2ZGt-COfekmJddQWXuCQAj7uY5YKcE3qEt7IBcaj96MNbF8b2qxTNbLrzgXioIzl0SIw8Ws-YUOu5I3A__&Key-Pair-Id=K3RPWS32NSSJCE\n",
+      "Resolving cdn-lfs.hf.co (cdn-lfs.hf.co)... 108.158.20.116, 108.158.20.30, 108.158.20.84, ...\n",
+      "Connecting to cdn-lfs.hf.co (cdn-lfs.hf.co)|108.158.20.116|:443... connected.\n",
       "HTTP request sent, awaiting response... 200 OK\n",
       "Length: 4265437280 (4.0G) [binary/octet-stream]\n",
       "Saving to: ‘sd-v1-5-inpainting.ckpt’\n",
       "\n",
-      "sd-v1-5-inpainting. 100%[===================>]   3.97G   299MB/s    in 12s     \n",
+      "sd-v1-5-inpainting. 100%[===================>]   3.97G   366MB/s    in 12s     \n",
       "\n",
-      "2025-06-11 10:33:13 (341 MB/s) - ‘sd-v1-5-inpainting.ckpt’ saved [4265437280/4265437280]\n",
+      "2025-06-13 07:07:46 (353 MB/s) - ‘sd-v1-5-inpainting.ckpt’ saved [4265437280/4265437280]\n",
       "\n"
      ]
     }
@@ -141,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "f7450c55",
    "metadata": {},
    "outputs": [
@@ -169,7 +154,7 @@
     }
    ],
    "source": [
-    "!wget https://huggingface.co/zhengchong/CatVTON/resolve/main/vitonhd-16k-512/attention/model.safetensors?download=true "
+    "# !wget https://huggingface.co/zhengchong/CatVTON/resolve/main/vitonhd-16k-512/attention/model.safetensors?download=true "
    ]
   },
   {
@@ -333,22 +318,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 247,
+   "execution_count": 4,
    "id": "91ef7a4e",
    "metadata": {},
    "outputs": [
     {
-     "ename": "KeyError",
-     "evalue": "'_oh'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
-      "\u001b[0;32m/tmp/ipykernel_71/1017109895.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mempty_cache\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Release unused GPU memory\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mgc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m              \u001b[0;31m# Run Python garbage collector\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
-      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/IPython/core/displayhook.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, result)\u001b[0m\n\u001b[1;32m    261\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite_output_prompt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    262\u001b[0m             \u001b[0mformat_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmd_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompute_format_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 263\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate_user_ns\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    264\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfill_exec_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mformat_dict\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/IPython/core/displayhook.py\u001b[0m in \u001b[0;36mupdate_user_ns\u001b[0;34m(self, result)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    200\u001b[0m         \u001b[0;31m# Avoid recursive reference when displaying _oh/Out\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 201\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcache_size\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshell\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_ns\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'_oh'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    202\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshell\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_ns\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'_oh'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcache_size\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdo_full_cache\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    203\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcull_cache\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mKeyError\u001b[0m: '_oh'"
-     ]
+     "data": {
+      "text/plain": [
+       "0"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
@@ -361,7 +343,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 249,
+   "execution_count": 9,
    "id": "08f29055",
    "metadata": {},
    "outputs": [
@@ -369,7 +351,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "GPU memory used: 8.12 MB / 16269.25 MB\n"
+      "GPU memory used: 0.00 MB / 16269.25 MB\n"
      ]
     }
    ],
@@ -396,7 +378,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
    "id": "37335c1e",
    "metadata": {},
    "outputs": [],
@@ -416,7 +398,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "id": "35d98b83",
    "metadata": {},
    "outputs": [],
@@ -450,14 +432,231 @@
    "id": "d7ff094a",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "from torch.nn import functional as F\n",
+    "import torch\n",
+    "# from flash_attn import flash_attn_func\n",
+    "\n",
+    "class SkipAttnProcessor(torch.nn.Module):\n",
+    "    def __init__(self, *args, **kwargs) -> None:\n",
+    "        super().__init__()\n",
+    "\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        attn,\n",
+    "        hidden_states,\n",
+    "        encoder_hidden_states=None,\n",
+    "        attention_mask=None,\n",
+    "        temb=None,\n",
+    "    ):\n",
+    "        return hidden_states\n",
+    "\n",
+    "class AttnProcessor2_0(torch.nn.Module):\n",
+    "    r\"\"\"\n",
+    "    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        hidden_size=None,\n",
+    "        cross_attention_dim=None,\n",
+    "        **kwargs\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        if not hasattr(F, \"scaled_dot_product_attention\"):\n",
+    "            raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n",
+    "\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        attn,\n",
+    "        hidden_states,\n",
+    "        encoder_hidden_states=None,\n",
+    "        attention_mask=None,\n",
+    "        temb=None,\n",
+    "        *args,\n",
+    "        **kwargs,\n",
+    "    ):\n",
+    "        residual = hidden_states\n",
+    "\n",
+    "        if attn.spatial_norm is not None:\n",
+    "            hidden_states = attn.spatial_norm(hidden_states, temb)\n",
+    "\n",
+    "        input_ndim = hidden_states.ndim\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            batch_size, channel, height, width = hidden_states.shape\n",
+    "            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n",
+    "\n",
+    "        batch_size, sequence_length, _ = (\n",
+    "            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n",
+    "        )\n",
+    "\n",
+    "        if attention_mask is not None:\n",
+    "            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n",
+    "            # scaled_dot_product_attention expects attention_mask shape to be\n",
+    "            # (batch, heads, source_length, target_length)\n",
+    "            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n",
+    "\n",
+    "        if attn.group_norm is not None:\n",
+    "            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n",
+    "\n",
+    "        query = attn.to_q(hidden_states)\n",
+    "\n",
+    "        if encoder_hidden_states is None:\n",
+    "            encoder_hidden_states = hidden_states\n",
+    "        elif attn.norm_cross:\n",
+    "            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n",
+    "\n",
+    "        key = attn.to_k(encoder_hidden_states)\n",
+    "        value = attn.to_v(encoder_hidden_states)\n",
+    "\n",
+    "        inner_dim = key.shape[-1]\n",
+    "        head_dim = inner_dim // attn.heads\n",
+    "\n",
+    "        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "\n",
+    "        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n",
+    "\n",
+    "        # the output of sdp = (batch, num_heads, seq_len, head_dim)\n",
+    "        # TODO: add support for attn.scale when we move to Torch 2.1\n",
+    "        \n",
+    "        hidden_states = F.scaled_dot_product_attention(\n",
+    "            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n",
+    "        )\n",
+    "        # hidden_states = flash_attn_func(\n",
+    "        #     query, key, value, dropout_p=0.0, causal=False\n",
+    "        # )\n",
+    "\n",
+    "        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n",
+    "        hidden_states = hidden_states.to(query.dtype)\n",
+    "\n",
+    "        # linear proj\n",
+    "        hidden_states = attn.to_out[0](hidden_states)\n",
+    "        # dropout\n",
+    "        hidden_states = attn.to_out[1](hidden_states)\n",
+    "\n",
+    "        if input_ndim == 4:\n",
+    "            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n",
+    "\n",
+    "        if attn.residual_connection:\n",
+    "            hidden_states = hidden_states + residual\n",
+    "\n",
+    "        hidden_states = hidden_states / attn.rescale_output_factor\n",
+    "\n",
+    "        return hidden_states\n",
+    "   "
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6028381d",
+   "id": "84a7fa87",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import torch\n",
+    "\n",
+    "def init_adapter(unet, \n",
+    "                 cross_attn_cls=SkipAttnProcessor,\n",
+    "                 self_attn_cls=None,\n",
+    "                 cross_attn_dim=None, \n",
+    "                 **kwargs):\n",
+    "    if cross_attn_dim is None:\n",
+    "        cross_attn_dim = unet.config.cross_attention_dim\n",
+    "    attn_procs = {}\n",
+    "    for name in unet.attn_processors.keys():\n",
+    "        cross_attention_dim = None if name.endswith(\"attn1.processor\") else cross_attn_dim\n",
+    "        if name.startswith(\"mid_block\"):\n",
+    "            hidden_size = unet.config.block_out_channels[-1]\n",
+    "        elif name.startswith(\"up_blocks\"):\n",
+    "            block_id = int(name[len(\"up_blocks.\")])\n",
+    "            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]\n",
+    "        elif name.startswith(\"down_blocks\"):\n",
+    "            block_id = int(name[len(\"down_blocks.\")])\n",
+    "            hidden_size = unet.config.block_out_channels[block_id]\n",
+    "        if cross_attention_dim is None:\n",
+    "            if self_attn_cls is not None:\n",
+    "                attn_procs[name] = self_attn_cls(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)\n",
+    "            else:\n",
+    "                # retain the original attn processor\n",
+    "                attn_procs[name] = AttnProcessor2_0(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)\n",
+    "        else:\n",
+    "            attn_procs[name] = cross_attn_cls(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)\n",
+    "                                                    \n",
+    "    unet.set_attn_processor(attn_procs)\n",
+    "    adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())\n",
+    "    return adapter_modules\n",
+    "\n",
+    "def init_diffusion_model(diffusion_model_name_or_path, unet_class=None):\n",
+    "    from diffusers import AutoencoderKL\n",
+    "    from transformers import CLIPTextModel, CLIPTokenizer\n",
+    "\n",
+    "    text_encoder = CLIPTextModel.from_pretrained(diffusion_model_name_or_path, subfolder=\"text_encoder\")\n",
+    "    vae = AutoencoderKL.from_pretrained(diffusion_model_name_or_path, subfolder=\"vae\")\n",
+    "    tokenizer = CLIPTokenizer.from_pretrained(diffusion_model_name_or_path, subfolder=\"tokenizer\")\n",
+    "    try:\n",
+    "        unet_folder = os.path.join(diffusion_model_name_or_path, \"unet\")\n",
+    "        unet_configs = json.load(open(os.path.join(unet_folder, \"config.json\"), \"r\"))\n",
+    "        unet = unet_class(**unet_configs)\n",
+    "        unet.load_state_dict(torch.load(os.path.join(unet_folder, \"diffusion_pytorch_model.bin\"), map_location=\"cpu\"), strict=True)\n",
+    "    except:\n",
+    "        unet = None\n",
+    "    return text_encoder, vae, tokenizer, unet\n",
+    "\n",
+    "def attn_of_unet(unet):\n",
+    "    attn_blocks = torch.nn.ModuleList()\n",
+    "    for name, param in unet.named_modules():\n",
+    "        if \"attn1\" in name:\n",
+    "            attn_blocks.append(param)\n",
+    "    return attn_blocks\n",
+    "\n",
+    "def get_trainable_module(unet, trainable_module_name):\n",
+    "    if trainable_module_name == \"unet\":\n",
+    "        return unet\n",
+    "    elif trainable_module_name == \"transformer\":\n",
+    "        trainable_modules = torch.nn.ModuleList()\n",
+    "        for blocks in [unet.down_blocks, unet.mid_block, unet.up_blocks]:\n",
+    "            if hasattr(blocks, \"attentions\"):\n",
+    "                trainable_modules.append(blocks.attentions)\n",
+    "            else:\n",
+    "                for block in blocks:\n",
+    "                    if hasattr(block, \"attentions\"):\n",
+    "                        trainable_modules.append(block.attentions)\n",
+    "        return trainable_modules\n",
+    "    elif trainable_module_name == \"attention\":\n",
+    "        attn_blocks = torch.nn.ModuleList()\n",
+    "        for name, param in unet.named_modules():\n",
+    "            if \"attn1\" in name:\n",
+    "                attn_blocks.append(param)\n",
+    "        return attn_blocks\n",
+    "    else:\n",
+    "        raise ValueError(f\"Unknown trainable_module_name: {trainable_module_name}\")\n",
+    "\n",
+    "                \n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6028381d",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'model'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_662/1349749640.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtransformers\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mCLIPImageProcessor\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattn_processor\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mSkipAttnProcessor\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mget_trainable_module\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minit_adapter\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m from utils import (check_inputs, get_time_embedding, numpy_to_pil, prepare_image,\n",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'model'"
+     ]
+    }
+   ],
    "source": [
     "import inspect\n",
     "import os\n",
@@ -475,8 +674,6 @@
     "from huggingface_hub import snapshot_download\n",
     "from transformers import CLIPImageProcessor\n",
     "\n",
-    "from model.attn_processor import SkipAttnProcessor\n",
-    "from model.utils import get_trainable_module, init_adapter\n",
     "from utils import (check_inputs, get_time_embedding, numpy_to_pil, prepare_image,\n",
     "                   prepare_mask_image, resize_and_crop, resize_and_padding)\n",
     "from ddpm import DDPMSampler\n",
@@ -755,7 +952,7 @@
     "from diffusers.image_processor import VaeImageProcessor\n",
     "from tqdm import tqdm\n",
     "from PIL import Image, ImageFilter\n",
-    "import model\n",
+    "import load_model\n",
     "\n",
     "from utils import repaint, to_pil_image\n",
     "                    \n",
@@ -921,7 +1118,8 @@
     "        \"base_model_path\": \"booksforcharlie/stable-diffusion-inpainting\",\n",
     "        \"resume_path\": \"zhengchong/CatVTON\",\n",
     "        \"dataset_name\": \"vitonhd\",\n",
-    "        \"data_root_path\": \"/kaggle/input/viton-hd-dataset\",\n",
+    "        # \"data_root_path\": \"/kaggle/input/viton-hd-dataset\",\n",
+    "        \"data_root_path\": \"/kaggle/working/stable-diffusion/sample_dataset\",\n",
     "        \"output_dir\": \"./output\",\n",
     "        \"seed\": 555,\n",
     "        \"batch_size\": 2,\n",
@@ -936,10 +1134,11 @@
     "        \"dataloader_num_workers\": 4,\n",
     "        \"mixed_precision\": 'no',\n",
     "        \"concat_axis\": 'y',\n",
-    "        \"enable_condition_noise\": True\n",
+    "        \"enable_condition_noise\": True,\n",
+    "        \"is_train\": False\n",
     "    }\n",
     "\n",
-    "    models=model.preload_models_from_standard_weights(ckpt_path=\"sd-v1-5-inpainting.ckpt\", device=\"cuda\", finetune_weight_path=\"model.safetensors\")\n",
+    "    models=load_model.preload_models_from_standard_weights(ckpt_path=\"sd-v1-5-inpainting.ckpt\", device=\"cuda\", finetune_weights_path=\"/kaggle/working/stable-diffusion/checkpoints/checkpoint_epoch_10.pth\")\n",
     "\n",
     "    # Pipeline\n",
     "    pipeline = CatVTONPipeline(\n",
@@ -1795,18 +1994,6 @@
    "display_name": "Python 3 (ipykernel)",
    "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.11.11"
   }
  },
  "nbformat": 4,