fix bug

inference_utils.py

@@ -9,6 +9,17 @@ torch.cuda.manual_seed_all(seed)
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
+# SPIGA ckpt downloading always fails, so we download it manually and put it in the right place.
+import site
+from gdown import download
+
+user_site_packages_path = site.getusersitepackages()
+spiga_file_id = "1YrbScfMzrAAWMJQYgxdLZ9l57nmTdpQC"
+ckpt_path = os.path.join(user_site_packages_path, "spiga/models/weights/spiga_300wpublic.pt")
+if not os.path.exists(ckpt_path):
+    os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
+    download(id=spiga_file_id, output=ckpt_path)
+
 from PIL import Image
 from gdown import download_folder
 from facelib import FaceDetector
@@ -53,21 +64,21 @@ def concatenate_images(image_files, output_file):
 
 def init_pipeline():
     # Initialize the model
-    model_id  = "runwayml/stable-diffusion-v1-5"  # or your local sdv1-5 path
+    model_id = "runwayml/stable-diffusion-v1-5"  # or your local sdv1-5 path
     base_path = "./checkpoints/stablemakeup"
     folder_id = "1397t27GrUyLPnj17qVpKWGwg93EcaFfg"
     if not os.path.exists(base_path):
         download_folder(id=folder_id, output=base_path, quiet=False, use_cookies=False)
     makeup_encoder_path = base_path + "/pytorch_model.bin"
-    id_encoder_path     = base_path + "/pytorch_model_1.bin"
-    pose_encoder_path   = base_path + "/pytorch_model_2.bin"
-
-    Unet              = OriginalUNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to("cuda")
-    id_encoder        = ControlNetModel.from_unet(Unet)
-    pose_encoder      = ControlNetModel.from_unet(Unet)
-    makeup_encoder    = detail_encoder(Unet, "openai/clip-vit-large-patch14", "cuda", dtype=torch.float32)
-    id_state_dict     = torch.load(id_encoder_path)
-    pose_state_dict   = torch.load(pose_encoder_path)
+    id_encoder_path = base_path + "/pytorch_model_1.bin"
+    pose_encoder_path = base_path + "/pytorch_model_2.bin"
+
+    Unet = OriginalUNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to("cuda")
+    id_encoder = ControlNetModel.from_unet(Unet)
+    pose_encoder = ControlNetModel.from_unet(Unet)
+    makeup_encoder = detail_encoder(Unet, "openai/clip-vit-large-patch14", "cuda", dtype=torch.float32)
+    id_state_dict = torch.load(id_encoder_path)
+    pose_state_dict = torch.load(pose_encoder_path)
     makeup_state_dict = torch.load(makeup_encoder_path)
     id_encoder.load_state_dict(id_state_dict, strict=False)
     pose_encoder.load_state_dict(pose_state_dict, strict=False)
@@ -82,14 +93,16 @@ def init_pipeline():
     pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
     return pipe, makeup_encoder
 
+
 # Initialize the model
 pipeline, makeup_encoder = init_pipeline()
 
 
 def inference(id_image_pil, makeup_image_pil, guidance_scale=1.6, size=512):
-    id_image     = id_image_pil.resize((size, size))
+    id_image = id_image_pil.resize((size, size))
     makeup_image = makeup_image_pil.resize((size, size))
-    pose_image   = get_draw(id_image, size=size)
-    result_img   = makeup_encoder.generate(id_image=[id_image, pose_image], makeup_image=makeup_image, pipe=pipeline, guidance_scale=guidance_scale)
+    pose_image = get_draw(id_image, size=size)
+    result_img = makeup_encoder.generate(
+        id_image=[id_image, pose_image], makeup_image=makeup_image, pipe=pipeline, guidance_scale=guidance_scale
+    )
     return result_img
-

spiga_draw.py

@@ -7,17 +7,6 @@ from facelib import FaceDetector
 from spiga.inference.config import ModelConfig
 from spiga.inference.framework import SPIGAFramework
 
-
-
-# SPIGA ckpt downloading always fails, so we download it manually and put it in the right place.
-import site
-from gdown import download
-user_site_packages_path = site.getusersitepackages()
-spiga_file_id = "1YrbScfMzrAAWMJQYgxdLZ9l57nmTdpQC"
-ckpt_path = os.path.join(user_site_packages_path, "spiga/models/weights/spiga_300wpublic.pt")
-if not os.path.exists(ckpt_path):
-    os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
-    download(id=spiga_file_id, output=ckpt_path)
 processor = SPIGAFramework(ModelConfig("300wpublic"))
 
 def center_crop(image, size):