update

ootd/inference_ootd_dc.py

@@ -23,6 +23,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers import AutoProcessor, CLIPVisionModelWithProjection
 from transformers import CLIPTextModel, CLIPTokenizer
+import requests
 
 VIT_PATH = "openai/clip-vit-large-patch14"
 VAE_PATH = "levihsu/OOTDiffusion"
@@ -36,23 +37,32 @@ class OOTDiffusionDC:
 
         vae = AutoencoderKL.from_pretrained(
             VAE_PATH,
-            subfolder="vae",
+            subfolder="checkpoints/ootd/vae",
             torch_dtype=torch.float16,
         )
 
         unet_garm = UNetGarm2DConditionModel.from_pretrained(
             UNET_PATH,
-            subfolder="unet_garm",
+            subfolder="checkpoints/ootd/ootd_hd/checkpoint-36000/unet_garm",
             torch_dtype=torch.float16,
             use_safetensors=True,
         )
         unet_vton = UNetVton2DConditionModel.from_pretrained(
             UNET_PATH,
-            subfolder="unet_vton",
+            subfolder="checkpoints/ootd/ootd_hd/checkpoint-36000/unet_vton",
             torch_dtype=torch.float16,
             use_safetensors=True,
         )
 
+        #判断文件是否存在
+        filePath = "/home/user/app/checkpoints/ootd/text_encoder/pytorch_model.bin"
+        if os.path.exists(filePath) == False:
+            url = "https://huggingface.co/yangjoe/pytorch_model/resolve/main/pytorch_model.bin"
+            response = requests.get(url)
+            #下载该文件
+            with open("/home/user/app/checkpoints/ootd/text_encoder/pytorch_model.bin", "wb") as f:
+                f.write(response.content)
+
         self.pipe = OotdPipeline.from_pretrained(
             MODEL_PATH,
             unet_garm=unet_garm,
@@ -75,8 +85,8 @@ class OOTDiffusionDC:
             subfolder="tokenizer",
         )
         self.text_encoder = CLIPTextModel.from_pretrained(
-            MODEL_PATH,
-            subfolder="text_encoder",
+            'yangjoe/pytorch_model',
+            #subfolder="text_encoder",
         ).to(self.gpu_id)
 
 

ootd/inference_ootd_hd.py

@@ -54,11 +54,14 @@ class OOTDiffusionHD:
             torch_dtype=torch.float16,
             use_safetensors=True,
         )
-        url = "https://huggingface.co/yangjoe/pytorch_model/resolve/main/pytorch_model.bin"
-        response = requests.get(url)
-        #下载该文件
-        with open("/home/user/app/checkpoints/ootd/text_encoder/pytorch_model.bin", "wb") as f:
-            f.write(response.content)
+        #判断文件是否存在
+        filePath = "/home/user/app/checkpoints/ootd/text_encoder/pytorch_model.bin"
+        if os.path.exists(filePath) == False:
+            url = "https://huggingface.co/yangjoe/pytorch_model/resolve/main/pytorch_model.bin"
+            response = requests.get(url)
+            #下载该文件
+            with open("/home/user/app/checkpoints/ootd/text_encoder/pytorch_model.bin", "wb") as f:
+                f.write(response.content)
 
         self.pipe = OotdPipeline.from_pretrained(
             MODEL_PATH,