Update utils.py

utils.py

@@ -2,7 +2,7 @@ import os
 os.environ["HF_HOME"] = "/data/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/data/huggingface"
 os.makedirs("/data/huggingface/hub", exist_ok=True)
-os.makedirs("/data/huggingface/clip_vision_model", exist_ok=True)
+# os.makedirs("/data/huggingface/clip_vision_model", exist_ok=True)
 
 import torch
 from diffusers import StableDiffusionImg2ImgPipeline
@@ -40,18 +40,18 @@ pipe.load_ip_adapter(
     weight_name=IPADAPTER_WEIGHT_NAME
 )
 
-# Load vision encoder and processor for IP-Adapter embedding
-vision_encoder = CLIPVisionModelWithProjection.from_pretrained(
-    "h94/IP-Adapter",             # repo_id (main IP-Adapter repo)
-    subfolder="clip_vision_model",# subfolder within the repo!
-    cache_dir=CACHE_DIR
-)
+# # Load vision encoder and processor for IP-Adapter embedding
+# vision_encoder = CLIPVisionModelWithProjection.from_pretrained(
+#     "h94/IP-Adapter",             # repo_id (main IP-Adapter repo)
+#     subfolder="clip_vision_model",# subfolder within the repo!
+#     cache_dir=CACHE_DIR
+# )
 
-image_processor = CLIPImageProcessor.from_pretrained(
-    "h94/IP-Adapter",
-    subfolder="clip_vision_model",
-    cache_dir=CACHE_DIR
-)
+# image_processor = CLIPImageProcessor.from_pretrained(
+#     "h94/IP-Adapter",
+#     subfolder="clip_vision_model",
+#     cache_dir=CACHE_DIR
+# )
 
 def generate_sticker(input_image: Image.Image, prompt: str):
     """
@@ -67,13 +67,16 @@ def generate_sticker(input_image: Image.Image, prompt: str):
 
     # Preprocess the image (resize, etc)
     face_img = input_image.convert("RGB").resize((224, 224))
-    inputs = image_processor(images=face_img, return_tensors="pt").to(DEVICE)
-    with torch.no_grad():
-        image_embeds = vision_encoder(**inputs).image_embeds
+    # inputs = image_processor(images=face_img, return_tensors="pt").to(DEVICE)
+    # with torch.no_grad():
+    #     image_embeds = vision_encoder(**inputs).image_embeds
 
     # 2. Prepare image for SD pipeline
     init_image = input_image.convert("RGB").resize((512, 512))
 
+    # IP-Adapter expects the reference image via image_embeds, which is produced by this function:
+    image_embeds = pipe.prepare_ip_adapter_image_embeds(face_img)
+
     # Run inference (low strength for identity preservation)
     result = pipe(
         prompt=prompt,