Update app.py

app.py

@@ -6,93 +6,128 @@ from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 import os
+import logging
 
-# Allow network access for runtime downloads
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Allow network access
 os.environ["HF_HUB_OFFLINE"] = "0"
 
-# Set device to CPU (Hugging Face free tier is CPU-only)
+# Set device to CPU
 device = "cpu"
-dtype = torch.float32  # Use float32 to avoid GPU-specific optimizations
+dtype = torch.float32
+
+# Define cache directory
+cache_dir = "./cache"
+os.makedirs(cache_dir, exist_ok=True)
 
-# Load face encoder (InsightFace will download its weights on first run)
+# Load face encoder
 try:
     face_app = FaceAnalysis(providers=["CPUExecutionProvider"])
     face_app.prepare(ctx_id=0, det_size=(480, 480))
-    print("InsightFace model loaded successfully.")
+    logger.info("InsightFace model loaded successfully.")
 except Exception as e:
-    raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access.")
+    raise RuntimeError(f"Failed to load InsightFace model: {e}")
 
-# Define paths for temporary storage (ephemeral in Spaces)
+# Download function with explicit path return
+def download_file(repo_id, filename, local_dir):
+    file_path = os.path.join(local_dir, filename)
+    if not os.path.exists(file_path):
+        logger.info(f"Downloading {filename} from {repo_id} to {local_dir}...")
+        try:
+            downloaded_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                local_dir=local_dir,
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
+            logger.info(f"Downloaded to {downloaded_path}")
+            return downloaded_path  # Return the actual path from hf_hub_download
+        except Exception as e:
+            logger.error(f"Download failed: {e}")
+            raise
+    else:
+        logger.info(f"Using cached file at {file_path}")
+        return file_path
+
+# Define paths
 kolors_unet_path = "./unet"
 ip_adapter_path = "./"
+os.makedirs(kolors_unet_path, exist_ok=True)
+os.makedirs(ip_adapter_path, exist_ok=True)
 
-# Download Kolors unet weights at runtime
-kolors_weights = os.path.join(kolors_unet_path, "diffusion_pytorch_model.fp16.safetensors")
-if not os.path.exists(kolors_weights):
-    print("Downloading Kolors unet weights...")
-    os.makedirs(kolors_unet_path, exist_ok=True)
-    hf_hub_download(
-        repo_id="Kwai-Kolors/Kolors",
-        filename="unet/diffusion_pytorch_model.fp16.safetensors",
-        local_dir=kolors_unet_path,
-        local_files_only=False
-    )
-    print("Kolors unet weights downloaded to", kolors_weights)
-
-# Download IP-Adapter weights at runtime
-ip_adapter_weights = os.path.join(ip_adapter_path, "ipa-faceid-plus.bin")
-if not os.path.exists(ip_adapter_weights):
-    print("Downloading IP-Adapter weights...")
-    hf_hub_download(
-        repo_id="Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
-        filename="ipa-faceid-plus.bin",
-        local_dir=ip_adapter_path,
-        local_files_only=False
-    )
-    print("IP-Adapter weights downloaded to", ip_adapter_weights)
-
-# Load the base SDXL pipeline directly from Hugging Face Hub
-print("Loading Stable Diffusion XL base model...")
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=dtype,
-    safety_checker=None,
-    local_files_only=False,  # Download from Hub at runtime
-    cache_dir="./cache"  # Use temporary cache directory
+# Download weights and get exact paths
+kolors_weights = download_file(
+    "Kwai-Kolors/Kolors",
+    "unet/diffusion_pytorch_model.fp16.safetensors",
+    kolors_unet_path
+)
+ip_adapter_weights = download_file(
+    "Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
+    "ipa-faceid-plus.bin",
+    ip_adapter_path
 )
 
-# Replace unet with Kolors weights
-print("Loading Kolors unet weights into pipeline...")
-pipe.unet.load_state_dict(torch.load(kolors_weights, map_location=device))
+# Load the pipeline
+logger.info("Loading Stable Diffusion XL base model...")
+try:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        torch_dtype=dtype,
+        safety_checker=None,
+        local_files_only=False,
+        cache_dir=cache_dir
+    )
+except Exception as e:
+    logger.error(f"Failed to load SDXL base model: {e}")
+    raise
+
+# Load Kolors unet weights
+logger.info(f"Loading Kolors unet weights from {kolors_weights}...")
+try:
+    state_dict = torch.load(kolors_weights, map_location=device)
+    pipe.unet.load_state_dict(state_dict)
+    logger.info("Kolors unet weights loaded successfully.")
+except Exception as e:
+    logger.error(f"Failed to load Kolors unet weights: {e}")
+    raise
 
 # Load IP-Adapter
-print("Loading IP-Adapter...")
-pipe.load_ip_adapter(ip_adapter_path, subfolder=None, weight_name="ipa-faceid-plus.bin")
+logger.info(f"Loading IP-Adapter from {ip_adapter_weights}...")
+try:
+    pipe.load_ip_adapter(ip_adapter_path, subfolder=None, weight_name="ipa-faceid-plus.bin")
+except Exception as e:
+    logger.error(f"Failed to load IP-Adapter: {e}")
+    raise
 
 # Move pipeline to CPU
 pipe.to(device)
 
 def generate_image(uploaded_image, prompt):
-    img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
-    faces = face_app.get(img)
-    if not faces:
-        return "No face detected!", None
+    try:
+        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]
-    face_emb = face_info["embedding"]
+        face_info = faces[-1]
+        face_emb = face_info["embedding"]
 
-    try:
-        # Reduce inference steps and resolution to fit free tier limits
+        logger.info(f"Generating image with prompt: {prompt}")
         image = pipe(
             prompt=prompt,
             image_embeds=face_emb,
-            num_inference_steps=15,  # Lower steps for faster execution
+            num_inference_steps=15,
             guidance_scale=7.5,
-            height=384,  # Smaller resolution to reduce memory usage
+            height=384,
             width=384,
         ).images[0]
         return "Image generated successfully!", image
     except Exception as e:
+        logger.error(f"Generation failed: {e}")
         return f"Generation failed: {e}", None
 
 # Gradio interface
@@ -106,8 +141,8 @@ interface = gr.Interface(
         gr.Textbox(label="Status"),
         gr.Image(label="Generated Image")
     ],
-    title="Face Reference Image Generator (Kolors-IP-Adapter-FaceID-Plus)",
-    description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
+    title="Face Reference Image Generator",
+    description="Upload an image with a face and generate a new image."
 )
 
 interface.launch()