Update app.py

app.py

@@ -1,146 +1,77 @@
 import gradio as gr
+import cv2
+import insightface
 import numpy as np
-import random
-from diffusers import DiffusionPipeline
-import torch
+import gradio.components as gr_comp
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load the insightface model once
+providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cv2.cuda.getCudaEnabledDeviceCount() > 0 else ["CPUExecutionProvider"]
+model_path = '/content/inswapper_128.onnx'
+model_swap_insightface = insightface.model_zoo.get_model(model_path, providers=providers)
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+# Prepare the FaceAnalysis model once
+FACE_ANALYSER = insightface.app.FaceAnalysis(
+    name="buffalo_l",
+    root=".", providers=providers, allowed_modules=["landmark_3d_68", "landmark_2d_106", "detection", "recognition"]
+)
+FACE_ANALYSER.prepare(ctx_id=0 if cv2.cuda.getCudaEnabledDeviceCount() > 0 else -1, det_size=(768, 512))
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+# Function to update template choices based on gender
+def update_templates(gender):
+    if gender == 'Male':
+        return gr.Dropdown.update(choices=['Boy Template 1.JPG', 'Boy Template 2.JPG', 'Boy Template 3.JPG'])
+    elif gender == 'Female':
+        return gr.Dropdown.update(choices=['Girl Template 1.JPG', 'Girl Template 2.JPG', 'Girl Template 3.JPG'])
+    else:
+        return gr.Dropdown.update(choices=[])
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+# Main function for face swapping
+def face_swap_and_merge(src_image, gender, template_choice):
+    # Resize the source image to 700x400 for uniform processing
+    src_image = cv2.resize(src_image, (400, 600))
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image
+    template_image_path = f'/content/{template_choice}'
+    template_image = cv2.imread(template_image_path, cv2.IMREAD_UNCHANGED)
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+    # Resize the template image to 700x400 for uniform processing
+    template_image = cv2.resize(template_image, (400, 600))
 
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
-"""
+    src_faces = FACE_ANALYSER.get(src_image)
+    template_faces = FACE_ANALYSER.get(template_image)
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
+    img_fake = model_swap_insightface.get(img=template_image, target_face=template_faces[0], source_face=src_faces[0], paste_back=True)
+    img_fake_rgb = cv2.cvtColor(img_fake, cv2.COLOR_BGR2RGB)
 
-with gr.Blocks(css=css) as demo:
-    
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
-        """)
-        
-        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)
+    return img_fake_rgb
 
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=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,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
+# Function to clear only the output image
+def clear_output():
+    return None
 
-    run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
-    )
+# Function to clear all inputs and outputs
+def clear_all():
+    return None, None, None, None
+
+# Create Gradio interface with title and theme
+with gr.Blocks(theme='upsatwal/mlsc_tiet') as iface:
+    gr.Markdown("# Face Changer")
+    gr.Markdown(" ")
+    
+    src_image = gr.Image(type="numpy", label="Input Image")
+    gender = gr.Dropdown(["Male", "Female"], label="Select Gender", interactive=True)
+    template_choice = gr.Dropdown([], label="Select Template", interactive=True)
+    
+    gender.change(fn=update_templates, inputs=gender, outputs=template_choice)
+    
+    result_image = gr.Image(label="Output Image")
+    
+    submit_button = gr.Button("Submit")
+    submit_button.click(fn=face_swap_and_merge, inputs=[src_image, gender, template_choice], outputs=result_image)
 
-demo.queue().launch()
+    clear_output_button = gr.Button("Clear Output")
+    clear_output_button.click(fn=clear_output, inputs=[], outputs=result_image)
+    
+    clear_all_button = gr.Button("Clear All")
+    clear_all_button.click(fn=clear_all, inputs=[], outputs=[src_image, gender, template_choice, result_image])
+    
+    iface.launch(debug=True)