Update app.py

app.py

@@ -31,8 +31,7 @@ if not os.path.exists(buffalo_l_zip):
         print("Download completed.")
     except Exception as e:
         print(f"Failed to download buffalo_l.zip: {e}")
-        # Fallback to a lighter model or raise error
-        raise RuntimeError("Cannot download buffalo_l.zip. Please ensure network access or preload the file.")
+        raise RuntimeError("Cannot download buffalo_l.zip. Please ensure network access.")
 
 # Extract buffalo_l.zip
 if os.path.exists(buffalo_l_zip):
@@ -46,12 +45,26 @@ face_app = FaceAnalysis(providers=["CPUExecutionProvider"], root="/home/user/.in
 face_app.prepare(ctx_id=0, det_size=(480, 480))
 
 # Define paths for preloaded weights
-model_path = "./"  # Kolors base model weights
+model_path = "./unet/"  # Adjusted for preloaded Kolors 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}")
+# Check if Kolors weights exist, download if missing
+kolors_weights = model_path + "diffusion_pytorch_model.safetensors"
+if not os.path.exists(model_path):
+    os.makedirs(model_path, exist_ok=True)
+if not os.path.exists(kolors_weights):
+    try:
+        print("Downloading Kolors model weights...")
+        urllib.request.urlretrieve(
+            "https://huggingface.co/Kwai-Kolors/Kolors-diffusers/resolve/main/unet/diffusion_pytorch_model.safetensors",
+            kolors_weights
+        )
+        print("Kolors weights downloaded.")
+    except Exception as e:
+        print(f"Failed to download Kolors weights: {e}")
+        raise RuntimeError("Cannot download Kolors weights. Please ensure network access.")
+
+# Check if IP-Adapter weights exist
 if not os.path.exists(ip_adapter_path + "ip-adapter.bin"):
     raise FileNotFoundError(f"IP-Adapter weights not found at {ip_adapter_path}")
 
@@ -82,19 +95,4 @@ def generate_image(uploaded_image, prompt):
             image_embeds=face_emb,
             num_inference_steps=20,
             guidance_scale=7.5,
-            height=512,
-            width=512,
-        ).images[0]
-        return "Image generated successfully!", image
-    except Exception as e:
-        return f"Generation failed: {e}", None
-
-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")],
-    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."
-)
-
-interface.launch()
+            height=512,