Update app.py

app.py

@@ -1,162 +1,51 @@
-import cv2
-import torch
-import numpy as np
 import gradio as gr
-from diffusers import StableDiffusionPipeline  # Use SD 2.1 instead of SDXL
-from insightface.app import FaceAnalysis
-from huggingface_hub import hf_hub_download
-import os
-import logging
-import time
-
-# Set up detailed logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Allow network access
-os.environ["HF_HUB_OFFLINE"] = "0"
-
-# Set device to CPU
-device = "cpu"
-dtype = torch.float32
-
-# Define cache directory
-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:
-    logger.error(f"Failed to load InsightFace model: {e}")
-    raise
-
-# Download function with retry logic
-def download_file(repo_id, filename, local_dir, max_retries=3):
-    file_path = os.path.join(local_dir, filename)
-    if not os.path.exists(file_path):
-        for attempt in range(max_retries):
-            logger.info(f"Attempt {attempt + 1}/{max_retries}: 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
-            except Exception as e:
-                logger.error(f"Download attempt {attempt + 1} failed: {e}")
-                if attempt < max_retries - 1:
-                    logger.info("Retrying in 5 seconds...")
-                    time.sleep(5)
-                else:
-                    raise RuntimeError(f"Failed to download {filename} after {max_retries} attempts: {e}")
-    else:
-        logger.info(f"Using cached file at {file_path}")
-        return file_path
-
-# Define paths
-ip_adapter_path = "./"
-os.makedirs(ip_adapter_path, exist_ok=True)
-
-# Download IP-Adapter weights with retries
-logger.info("Starting weights download...")
-ip_adapter_weights = download_file(
-    "Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",
-    "ipa-faceid-plus.bin",
-    ip_adapter_path
-)
-
-# Load the pipeline with SD 2.1
-logger.info("Loading Stable Diffusion 2.1 base model...")
-try:
-    max_retries = 3
-    for attempt in range(max_retries):
-        try:
-            logger.info(f"Attempt {attempt + 1}/{max_retries}: Loading SD 2.1 model...")
-            pipe = StableDiffusionPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-2-1",
-                torch_dtype=dtype,
-                safety_checker=None,
-                local_files_only=False,
-                cache_dir=cache_dir,
-                variant="fp16",
-                use_safetensors=True
-            )
-            logger.info("SD 2.1 base model loaded successfully.")
-            break
-        except Exception as e:
-            logger.error(f"Load attempt {attempt + 1} failed: {e}")
-            if attempt < max_retries - 1:
-                logger.info("Retrying in 5 seconds...")
-                time.sleep(5)
-            else:
-                raise RuntimeError(f"Failed to load SD 2.1 model after {max_retries} attempts: {e}")
-except Exception as e:
-    logger.error(f"Failed to load SD 2.1 base model: {e}")
-    raise
-
-# Load IP-Adapter
-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]
-        face_emb = face_info["embedding"]
-
-        logger.info(f"Generating image with prompt: {prompt}")
-        image = pipe(
-            prompt=prompt,
-            image_embeds=face_emb,
-            num_inference_steps=10,
-            guidance_scale=7.5,
-            height=256,
-            width=256
-        ).images[0]
-        logger.info("Image generated successfully.")
-        return "Image generated successfully!", image
-    except Exception as e:
-        logger.error(f"Generation failed: {e}")
-        return f"Generation failed: {e}", None
-
-# Gradio interface
+from diffusers import StableDiffusionImg2ImgPipeline
+import torch
+from PIL import Image
+from codeformer_app import CodeFormerFaceRestoration
+
+# Load models
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    use_safetensors=True
+).to("cuda")
+
+codeformer = CodeFormerFaceRestoration()
+
+# Define the image-to-image function
+def generate_image(input_image, prompt, strength, fidelity):
+    # Preprocess the input image
+    init_image = Image.fromarray(input_image).convert("RGB")
+    init_image = init_image.resize((512, 512))
+
+    # Generate the image
+    generated_image = pipe(
+        prompt=prompt,
+        image=init_image,
+        strength=strength,
+        guidance_scale=7.5,
+        num_inference_steps=50
+    ).images[0]
+
+    # Restore the face
+    restored_image = codeformer.restore(generated_image, fidelity=fidelity)
+
+    return restored_image
+
+# Create the 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")
+        gr.Image(label="Upload Your Image"),  # Image upload input
+        gr.Textbox(label="Prompt"),           # Text input for the prompt
+        gr.Slider(0.1, 1.0, value=0.5, label="Strength (Lower = More Preservation)"),  # Strength slider
+        gr.Slider(0.1, 1.0, value=0.8, label="Fidelity (Higher = More Preservation)")  # Fidelity slider
     ],
-    title="Face Reference Image Generator",
-    description="Upload an image with a face and generate a new image."
+    outputs=gr.Image(label="Generated Image"),  # Output image
+    title="Image-to-Image with Face Preservation",
+    description="Upload an image, enter a prompt, and generate a new image while preserving the face."
 )
 
-logger.info("Launching Gradio interface...")
+# Launch the app
 interface.launch()