Update app.py

app.py

@@ -6,6 +6,7 @@ from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import os
+import zipfile
 
 # Force offline mode
 os.environ["HF_HUB_OFFLINE"] = "1"
@@ -14,19 +15,26 @@ os.environ["HF_HUB_OFFLINE"] = "1"
 device = "cpu"
 dtype = torch.float32
 
-# Load face encoder
-face_app = FaceAnalysis(providers=["CPUExecutionProvider"])
+# Extract InsightFace model if needed
+insightface_model_dir = "/home/user/.insightface/models/buffalo_l"
+if not os.path.exists(insightface_model_dir):
+    os.makedirs(insightface_model_dir, exist_ok=True)
+    with zipfile.ZipFile("./buffalo_l.zip", "r") as zip_ref:
+        zip_ref.extractall(insightface_model_dir)
+
+# Load face encoder with preloaded model
+face_app = FaceAnalysis(providers=["CPUExecutionProvider"], root="/home/user/.insightface/models")
 face_app.prepare(ctx_id=0, det_size=(480, 480))
 
 # Define paths for preloaded weights
-model_path = "./"  # Kolors base model
-ip_adapter_path = "./ip-adapter-faceid-plus-sdxl"
+model_path = "./"  # Kolors base model weights
+ip_adapter_path = "./"
 
 # Check if files exist
 if not os.path.exists(model_path + "diffusion_pytorch_model.safetensors"):
     raise FileNotFoundError(f"Kolors model weights not found at {model_path}")
-if not os.path.exists(ip_adapter_path) or not os.path.exists(os.path.join(ip_adapter_path, "config.json")):
-    raise FileNotFoundError(f"IP-Adapter weights or config.json not found at {ip_adapter_path}")
+if not os.path.exists(ip_adapter_path + "ip-adapter.bin"):
+    raise FileNotFoundError(f"IP-Adapter weights not found at {ip_adapter_path}")
 
 # Initialize model with empty weights
 with init_empty_weights():
@@ -41,13 +49,12 @@ pipe = load_checkpoint_and_dispatch(pipe, model_path, device_map="cpu", offload_
 pipe.load_ip_adapter("h94/IP-Adapter-FaceID-Plus-SDXL", subfolder=None, weight_name="ip-adapter.bin")
 
 def generate_image(uploaded_image, prompt):
-    # Convert Gradio image to OpenCV format
     img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
     faces = face_app.get(img)
     if not faces:
         return "No face detected!", None
 
-    face_info = faces[-1]  # Use largest face
+    face_info = faces[-1]
     face_emb = face_info["embedding"]
 
     try:
@@ -63,17 +70,10 @@ def generate_image(uploaded_image, prompt):
     except Exception as e:
         return f"Generation failed: {e}", None
 
-# Gradio interface
 interface = gr.Interface(
     fn=generate_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Reference Image"),
-        gr.Textbox(label="Enter Prompt", placeholder="e.g., A photorealistic astronaut in space")
-    ],
-    outputs=[
-        gr.Textbox(label="Status"),
-        gr.Image(label="Generated Image")
-    ],
+    inputs=[gr.Image(type="pil", label="Upload Reference Image"), gr.Textbox(label="Enter Prompt", placeholder="e.g., A photorealistic astronaut in space")],
+    outputs=[gr.Textbox(label="Status"), gr.Image(label="Generated Image")],
     title="Face Reference Image Generator (Kolors with IP-Adapter)",
     description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
 )