Update app.py

app.py

@@ -4,102 +4,73 @@ import numpy as np
 import gradio as gr
 from diffusers import StableDiffusionXLPipeline
 from insightface.app import FaceAnalysis
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
-import os
 from huggingface_hub import hf_hub_download
-import time
+import os
 
-# Allow network access for runtime download
+# Allow network access for runtime downloads
 os.environ["HF_HUB_OFFLINE"] = "0"
 
-# Set device to CPU
+# Set device to CPU (Hugging Face free tier is CPU-only)
 device = "cpu"
-dtype = torch.float32
+dtype = torch.float32  # Use float32 to avoid GPU-specific optimizations
 
-# Load face encoder (InsightFace handles its own download)
+# Load face encoder (InsightFace will download its weights on first run)
 try:
     face_app = FaceAnalysis(providers=["CPUExecutionProvider"])
     face_app.prepare(ctx_id=0, det_size=(480, 480))
     print("InsightFace model loaded successfully.")
 except Exception as e:
-    raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access for initial download.")
+    raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access.")
 
-# Define paths for preloaded or downloaded weights
-model_path = "./"  # Start with root
+# Define paths for temporary storage (ephemeral in Spaces)
+kolors_unet_path = "./unet"
 ip_adapter_path = "./"
 
-# Debug: List files to confirm preloading or download
-print("Files in root directory:", os.listdir("."))
-print("Files in ./unet/ directory:", os.listdir("./unet") if os.path.exists("./unet") else "No ./unet/ directory")
-
-# Check if base model weights exist or download them
-kolors_weights = model_path + "diffusers_weights.safetensors"
+# 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):
-    kolors_weights = model_path + "diffusion_pytorch_model.fp16.safetensors"
-    if not os.path.exists(kolors_weights):
-        kolors_weights_unet = "./unet/diffusion_pytorch_model.fp16.safetensors"
-        if not os.path.exists(kolors_weights_unet):
-            print("Preloading failed. Attempting runtime download with retry...")
-            os.makedirs("./unet", exist_ok=True)
-            max_retries = 3
-            for attempt in range(max_retries):
-                try:
-                    print(f"Download attempt {attempt + 1} of {max_retries}")
-                    hf_hub_download(
-                        repo_id="Kwai-Kolors/Kolors",
-                        filename="unet/diffusion_pytorch_model.fp16.safetensors",
-                        local_dir="./unet",
-                        local_files_only=False
-                    )
-                    print("Kolors base weights downloaded to", kolors_weights_unet)
-                    model_path = "./unet/"
-                    kolors_weights = kolors_weights_unet
-                    break
-                except Exception as e:
-                    print(f"Download attempt {attempt + 1} failed: {e}")
-                    if attempt < max_retries - 1:
-                        time.sleep(5)
-                    else:
-                        raise FileNotFoundError(f"Failed to download Kolors base weights after {max_retries} attempts: {e}. Verify the repo or contact support.")
-        else:
-            model_path = "./unet/"
-            kolors_weights = kolors_weights_unet
+    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)
 
-# Check if IP-Adapter weights exist or download them
-ip_adapter_weights = ip_adapter_path + "ipa-faceid-plus.bin"
+# 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("IP-Adapter preloading failed. Attempting runtime download with retry...")
-    max_retries = 3
-    for attempt in range(max_retries):
-        try:
-            print(f"IP-Adapter download attempt {attempt + 1} of {max_retries}")
-            hf_hub_download(
-                repo_id="Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
-                filename="ipa-faceid-plus.bin",
-                local_dir="./",
-                local_files_only=False
-            )
-            print("IP-Adapter weights downloaded to", ip_adapter_weights)
-            break
-        except Exception as e:
-            print(f"IP-Adapter download attempt {attempt + 1} failed: {e}")
-            if attempt < max_retries - 1:
-                time.sleep(5)
-            else:
-                raise FileNotFoundError(f"Failed to download IP-Adapter weights after {max_retries} attempts: {e}. Verify the repo or contact support.")
-
-# Initialize model with empty weights
-with init_empty_weights():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "./",  # Use local model path
-        torch_dtype=dtype,
-        safety_checker=None,
-        local_files_only=True  # Force local file usage
+    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
+)
+
+# 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 IP-Adapter
+print("Loading IP-Adapter...")
+pipe.load_ip_adapter(ip_adapter_path, subfolder=None, weight_name="ipa-faceid-plus.bin")
 
-# Load and dispatch model with accelerate
-pipe = load_checkpoint_and_dispatch(pipe, model_path, device_map="cpu", offload_folder=None)
-pipe.load_ip_adapter("./", subfolder=None, weight_name="ipa-faceid-plus.bin")
+# Move pipeline to CPU
+pipe.to(device)
 
 def generate_image(uploaded_image, prompt):
     img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
@@ -111,24 +82,32 @@ def generate_image(uploaded_image, prompt):
     face_emb = face_info["embedding"]
 
     try:
+        # Reduce inference steps and resolution to fit free tier limits
         image = pipe(
             prompt=prompt,
             image_embeds=face_emb,
-            num_inference_steps=20,
+            num_inference_steps=15,  # Lower steps for faster execution
             guidance_scale=7.5,
-            height=512,
-            width=512,
+            height=384,  # Smaller resolution to reduce memory usage
+            width=384,
         ).images[0]
         return "Image generated successfully!", image
     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-IP-Adapter-FaceID-Plus)",
     description="Upload an image with a face, enter a prompt, and generate a new image preserving the reference face."
 )
 
-interface.launch()
+interface.launch()