Update app.py

app.py

@@ -5,97 +5,51 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline, DDIMScheduler, IPAdapterFaceIDPipeline
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
-import os
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 device = "cpu"
 dtype = torch.float32
 
-# Directories
-cache_dir = "./cache"
-os.makedirs(cache_dir, exist_ok=True)
-ip_adapter_dir = "./ip_adapter"
-os.makedirs(ip_adapter_dir, exist_ok=True)
+# Initialize Face Detection
+face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
+face_app.prepare(ctx_id=0, det_size=(320, 320))
 
-# Download IP-Adapter FaceID weights
-logger.info("Downloading IP-Adapter FaceID weights...")
+# Download IP-Adapter FaceID Weights
 ip_adapter_path = hf_hub_download(
     repo_id="h94/IP-Adapter",
     filename="ip-adapter_sd15_faceid.bin",
-    subfolder="models",
-    local_dir=ip_adapter_dir,
-    cache_dir=cache_dir
+    subfolder="models"
 )
-logger.info(f"IP-Adapter FaceID weights downloaded to {ip_adapter_path}")
-
-# Initialize InsightFace
-logger.info("Initializing InsightFace...")
-face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
-face_app.prepare(ctx_id=0, det_size=(320, 320))
-logger.info("InsightFace initialized successfully.")
 
-# Load Stable Diffusion with IP-Adapter FaceID
-logger.info("Loading Stable Diffusion with IP-Adapter FaceID...")
-pipe = IPAdapterFaceIDPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=dtype,
-    cache_dir=cache_dir,
-    safety_checker=None
-)
+# Load Pipeline
+pipe = IPAdapterFaceIDPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=dtype)
 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 pipe.load_ip_adapter(ip_adapter_path)
 pipe.to(device)
-logger.info("Stable Diffusion with IP-Adapter FaceID loaded successfully.")
-
-# Generation function
-def generate_image(uploaded_image, prompt):
-    logger.info("Starting image generation...")
-    if uploaded_image is None:
-        return "Please upload an image!", None
 
-    logger.info("Processing uploaded image...")
-    img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
+def generate(upload_image, prompt):
+    if upload_image is None:
+        return "No Image Uploaded", None
+    
+    img = cv2.cvtColor(np.array(upload_image), cv2.COLOR_RGB2BGR)
     faces = face_app.get(img)
-    if not faces:
-        logger.warning("No face detected in the uploaded image.")
-        return "No face detected!", None
 
-    logger.info("Extracting face embedding...")
-    face_info = faces[0]
-    face_emb = torch.from_numpy(face_info.normed_embedding).unsqueeze(0).to(device)
-
-    logger.info(f"Generating image with prompt: {prompt}")
-    image = pipe(
-        prompt=prompt,
-        face_embeds=face_emb,
-        num_inference_steps=30,
-        guidance_scale=7.5,
-        height=512,
-        width=512,
-        generator=torch.Generator(device="cpu").manual_seed(42)
-    ).images[0]
+    if len(faces) == 0:
+        return "No Face Detected", None
     
-    logger.info("Image generated successfully.")
-    return "Image generated successfully!", image
+    face_emb = torch.tensor(faces[0].normed_embedding).unsqueeze(0).to(device)
+    result = pipe(prompt=prompt, face_embeds=face_emb, num_inference_steps=30).images[0]
+    return "Image Generated Successfully", result
 
-# Gradio interface
-interface = gr.Interface(
-    fn=generate_image,
+demo = gr.Interface(
+    generate,
     inputs=[
-        gr.Image(type="pil", label="Upload Reference Image"),
-        gr.Textbox(label="Prompt", placeholder="e.g., A man in red suit with sunglasses")
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Prompt")
     ],
     outputs=[
-        gr.Textbox(label="Status", interactive=False),
+        gr.Textbox(label="Status"),
         gr.Image(label="Generated Image")
-    ],
-    title="Face Reference Image Generator",
-    description="Upload an image with a face and generate a new image with the face preserved."
+    ]
 )
 
-logger.info("Launching Gradio interface...")
-interface.launch()
+demo.launch()