Update app.py

app.py

@@ -8,7 +8,7 @@ from huggingface_hub import hf_hub_download
 import os
 import logging
 
-# Set up logging
+# Set up detailed logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -24,12 +24,14 @@ cache_dir = "./cache"
 os.makedirs(cache_dir, exist_ok=True)
 
 # Load face encoder
+logger.info("Starting InsightFace initialization...")
 try:
     face_app = FaceAnalysis(providers=["CPUExecutionProvider"])
     face_app.prepare(ctx_id=0, det_size=(480, 480))
     logger.info("InsightFace model loaded successfully.")
 except Exception as e:
-    raise RuntimeError(f"Failed to load InsightFace model: {e}")
+    logger.error(f"Failed to load InsightFace model: {e}")
+    raise
 
 # Download function with explicit path return
 def download_file(repo_id, filename, local_dir):
@@ -59,7 +61,8 @@ ip_adapter_path = "./"
 os.makedirs(kolors_unet_path, exist_ok=True)
 os.makedirs(ip_adapter_path, exist_ok=True)
 
-# Download weights and get exact paths
+# Download weights
+logger.info("Starting weights download...")
 kolors_weights = download_file(
     "Kwai-Kolors/Kolors",
     "unet/diffusion_pytorch_model.fp16.safetensors",
@@ -71,7 +74,7 @@ ip_adapter_weights = download_file(
     ip_adapter_path
 )
 
-# Load the pipeline
+# Load the pipeline with verbose logging
 logger.info("Loading Stable Diffusion XL base model...")
 try:
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -79,13 +82,16 @@ try:
         torch_dtype=dtype,
         safety_checker=None,
         local_files_only=False,
-        cache_dir=cache_dir
+        cache_dir=cache_dir,
+        variant="fp16",  # Use FP16 weights to reduce memory usage
+        use_safetensors=True  # Prefer safetensors format if available
     )
+    logger.info("SDXL base model loaded successfully.")
 except Exception as e:
     logger.error(f"Failed to load SDXL base model: {e}")
     raise
 
-# Load Kolors unet weights with weights_only=False
+# Load Kolors unet weights
 logger.info(f"Loading Kolors unet weights from {kolors_weights}...")
 try:
     state_dict = torch.load(kolors_weights, map_location=device, weights_only=False)
@@ -99,18 +105,23 @@ except Exception as e:
 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")
+    logger.info("IP-Adapter loaded successfully.")
 except Exception as e:
     logger.error(f"Failed to load IP-Adapter: {e}")
     raise
 
 # Move pipeline to CPU
+logger.info("Moving pipeline to CPU...")
 pipe.to(device)
+logger.info("Pipeline moved to CPU.")
 
 def generate_image(uploaded_image, prompt):
+    logger.info("Starting image generation...")
     try:
         img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
         faces = face_app.get(img)
         if not faces:
+            logger.warning("No face detected in uploaded image.")
             return "No face detected!", None
 
         face_info = faces[-1]
@@ -120,11 +131,12 @@ def generate_image(uploaded_image, prompt):
         image = pipe(
             prompt=prompt,
             image_embeds=face_emb,
-            num_inference_steps=15,
+            num_inference_steps=10,  # Reduced steps for faster execution
             guidance_scale=7.5,
-            height=384,
-            width=384,
+            height=256,  # Smaller resolution to fit memory
+            width=256
         ).images[0]
+        logger.info("Image generated successfully.")
         return "Image generated successfully!", image
     except Exception as e:
         logger.error(f"Generation failed: {e}")
@@ -145,4 +157,5 @@ interface = gr.Interface(
     description="Upload an image with a face and generate a new image."
 )
 
+logger.info("Launching Gradio interface...")
 interface.launch()