Update app.py

app.py

@@ -1,67 +1,85 @@
 import gradio as gr
 import numpy as np
-import cv2
-import onnxruntime
-from insightface.app import FaceAnalysis
-from pathlib import Path
+import random
+import torch
+from diffusers import DiffusionPipeline
 
-# Initialize Face Analysis
-face_analyzer = FaceAnalysis(name="buffalo_l")
-face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
+# Ensure the model runs on CPU
+device = "cpu"
+dtype = torch.float32  # Use float32 for CPU compatibility
 
-# Load Face Swapper Model
-MODEL_PATH = Path("models/inswapper_128.onnx")
-if not MODEL_PATH.exists():
-    raise FileNotFoundError("Model file inswapper_128.onnx not found.")
+# Load model from Hugging Face (it will cache locally in Hugging Face Spaces)
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", 
+    torch_dtype=dtype, 
+    low_cpu_mem_usage=True
+).to(device)
 
-session = onnxruntime.InferenceSession(str(MODEL_PATH))
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
 
-def swap_faces(source_img, target_img):
-    """Perform face swapping using the ONNX model."""
-    try:
-        # Convert images to correct format
-        source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
-        target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
+def infer(prompt, seed=42, randomize_seed=False, width=512, height=512, num_inference_steps=4):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    image = pipe(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=0.0
+    ).images[0]
+    return image, seed
 
-        # Detect faces
-        source_faces = face_analyzer.get(source_img)
-        target_faces = face_analyzer.get(target_img)
+examples = [
+    "a tiny astronaut hatching from an egg on the moon",
+    "a cat holding a sign that says hello world",
+    "an anime illustration of a wiener schnitzel",
+]
 
-        if not source_faces or not target_faces:
-            return "No faces detected in one or both images."
-        if len(source_faces) > 1 or len(target_faces) > 1:
-            return "Multiple faces detected; only one face per image is supported."
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
 
-        source_face = source_faces[0]
-        target_face = target_faces[0]
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("""# FLUX.1 [schnell]
+        12B param rectified flow transformer distilled from FLUX.1 [pro]
+        """)
+        
+        with gr.Row():
+            prompt = gr.Text(label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False)
+            run_button = gr.Button("Run", scale=0)
+        
+        result = gr.Image(label="Result", show_label=False)
+        
+        with gr.Accordion("Advanced Settings", open=False):
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            
+            with gr.Row():
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)
+            
+            num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4)
+        
+        gr.Examples(
+            examples=examples,
+            fn=infer,
+            inputs=[prompt],
+            outputs=[result, seed],
+            cache_examples="lazy"
+        )
+    
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
+        outputs=[result, seed]
+    )
 
-        # Prepare input data for ONNX model
-        input_data = {
-            "target_image": target_img,
-            "target_face": target_face.embedding,
-            "source_face": source_face.embedding
-        }
-
-        # Run the ONNX model
-        result = session.run(None, input_data)[0]
-
-        # Convert result to image format
-        result_img = np.clip(result * 255, 0, 255).astype(np.uint8)
-        result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
-
-        return result_img
-    except Exception as e:
-        return f"Face swap failed: {e}"
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# Face Swap Tool 🚀")
-    with gr.Row():
-        input_source = gr.Image(label="Source Face", type="pil")
-        input_target = gr.Image(label="Target Image", type="pil")
-    btn_swap = gr.Button("Swap Faces")
-    output_image = gr.Image(label="Swapped Face")
-    btn_swap.click(swap_faces, inputs=[input_source, input_target], outputs=output_image)
-
-# Launch Gradio App
 demo.launch()