Update app.py

app.py

@@ -1,5 +1,6 @@
-# -* coding:UTF-8 -*
+# -* coding:UTF-8 *-
 # !/usr/bin/env python
+
 import numpy as np
 import gradio as gr
 import roop.globals
@@ -13,63 +14,50 @@ from roop.processors.frame.core import get_frame_processors_modules
 from roop.utilities import normalize_output_path
 import os
 from PIL import Image
-import cv2  # OpenCV for video processing
-import tempfile
-
-
-def save_video_frames(video_path, folder):
-    """Extract frames from video and save to folder"""
-    cap = cv2.VideoCapture(video_path)
-    frame_paths = []
-    idx = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame_path = os.path.join(folder, f"frame_{idx:05d}.png")
-        cv2.imwrite(frame_path, frame)
-        frame_paths.append(frame_path)
-        idx += 1
-    cap.release()
-    return frame_paths
-
-
-def swap_face(source_file, target_file, doFaceEnhancer):
-    # source_file and target_file are paths to video files
-
-    # Create temp folders to extract frames
-    temp_dir = tempfile.mkdtemp()
-    source_frames_dir = os.path.join(temp_dir, "source_frames")
-    target_frames_dir = os.path.join(temp_dir, "target_frames")
-    os.makedirs(source_frames_dir, exist_ok=True)
-    os.makedirs(target_frames_dir, exist_ok=True)
-
-    # Extract frames from source and target videos
-    source_frame_paths = save_video_frames(source_file, source_frames_dir)
-    target_frame_paths = save_video_frames(target_file, target_frames_dir)
-
-    # For demo, assume same number of frames (more advanced syncing may be needed)
-    # Save first frames as example input images for roop (or process each frame one by one)
-    source_img_path = os.path.join(temp_dir, "input.jpg")
-    target_img_path = os.path.join(temp_dir, "target.jpg")
-    first_source_frame = cv2.imread(source_frame_paths[0])
-    first_target_frame = cv2.imread(target_frame_paths[0])
-    cv2.imwrite(source_img_path, first_source_frame)
-    cv2.imwrite(target_img_path, first_target_frame)
-
-    roop.globals.source_path = source_img_path
-    roop.globals.target_path = target_img_path
-
-    output_path = os.path.join(temp_dir, "output_video.mp4")
+import mimetypes
+
+
+def swap_face(source_file, target_file, swap_type, doFaceEnhancer):
+    supported_images = ['.jpg', '.jpeg', '.png']
+    supported_videos = ['.mp4', '.mov', '.avi']
+
+    # Save source image
+    source_path = "input.jpg"
+    source_image = Image.fromarray(source_file)
+    source_image.save(source_path)
+
+    # Save target as image or video
+    mime_type, _ = mimetypes.guess_type(target_file.name)
+    ext = os.path.splitext(target_file.name)[1].lower()
+
+    if mime_type and "video" in mime_type:
+        target_path = "target" + ext
+        with open(target_path, "wb") as f:
+            f.write(target_file.read())
+    else:
+        target_path = "target.jpg"
+        target_image = Image.fromarray(target_file)
+        target_image.save(target_path)
+
+    print("source_path:", source_path)
+    print("target_path:", target_path)
+
+    roop.globals.source_path = source_path
+    roop.globals.target_path = target_path
+
+    # Output path based on type
+    if "video" in swap_type.lower() or ext in supported_videos:
+        output_path = "output.mp4"
+    else:
+        output_path = "output.jpg"
+
     roop.globals.output_path = normalize_output_path(
         roop.globals.source_path, roop.globals.target_path, output_path
     )
 
-    if doFaceEnhancer:
-        roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
-    else:
-        roop.globals.frame_processors = ["face_swapper"]
-
+    roop.globals.frame_processors = (
+        ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
+    )
     roop.globals.headless = True
     roop.globals.keep_fps = True
     roop.globals.keep_audio = True
@@ -81,25 +69,36 @@ def swap_face(source_file, target_file, doFaceEnhancer):
     roop.globals.execution_providers = decode_execution_providers(["cuda"])
     roop.globals.execution_threads = suggest_execution_threads()
 
-    # Check frame processors before start
+    print(
+        "start process",
+        roop.globals.source_path,
+        roop.globals.target_path,
+        roop.globals.output_path,
+    )
+
     for frame_processor in get_frame_processors_modules(
         roop.globals.frame_processors
     ):
         if not frame_processor.pre_check():
-            return None
+            return "Frame processor pre-check failed"
 
-    # Here you would ideally call start() and it processes the entire video,
-    # but since roop's start() expects images, you might need to adapt it
-    # to work frame by frame or modify roop internals accordingly.
     start()
-
-    # Return output video path to Gradio to play/download
-    return output_path
+    return roop.globals.output_path
 
 
 app = gr.Interface(
     fn=swap_face,
-    inputs=[gr.Video(label="Source Video"), gr.Video(label="Target Video"), gr.Checkbox(label="Face Enhancer?")],
-    outputs=gr.Video(label="Output Video"),
+    inputs=[
+        gr.Image(label="Source Image"),
+        gr.File(label="Target (Image or Video)"),
+        gr.Radio(choices=["Image to Image", "Image to Video"], label="Swap Type"),
+        gr.Checkbox(label="Face Enhancer?", info="Apply face enhancement")
+    ],
+    outputs=[
+        gr.Image(label="Output Image") if not os.path.exists("output.mp4") else gr.Video(label="Output Video")
+    ],
+    title="Face Swap Tool (Roop)",
+    description="Swap face from source image to a target image or video using Roop.",
 )
+
 app.launch()