Upload __init__.py

roop/__init__.py

@@ -1 +1,45 @@
-# This file contains the run_face_swap() function
+import os
+import shutil
+import cv2
+from moviepy.editor import VideoFileClip
+from PIL import Image
+from roop.core import run
+from tqdm import tqdm
+
+def run_face_swap(source_img, target_file):
+    filename = target_file.name
+
+    if filename.endswith(".jpg") or filename.endswith(".png"):
+        # Image swap (simplified)
+        os.makedirs("output", exist_ok=True)
+        output_path = "output/result.jpg"
+        os.system(f"python run.py --source {source_img} --target {filename} --output {output_path}")
+        return output_path
+
+    elif filename.endswith(".mp4"):
+        os.makedirs("frames_output", exist_ok=True)
+        os.makedirs("frames_input", exist_ok=True)
+        clip = VideoFileClip(filename)
+        fps = clip.fps
+        for i, frame in enumerate(clip.iter_frames()):
+            Image.fromarray(frame).save(f"frames_input/frame_{i:05d}.jpg")
+        import sys
+        sys.argv = ["run.py", "--source", source_img, "--target", "", "--output", ""]
+        for frame_file in sorted(os.listdir("frames_input")):
+            input_path = os.path.join("frames_input", frame_file)
+            output_path = os.path.join("frames_output", frame_file)
+            sys.argv[4] = input_path
+            sys.argv[6] = output_path
+            run()
+        frame_files = sorted(os.listdir("frames_output"))
+        frame_images = [cv2.imread(f"frames_output/{f}") for f in frame_files]
+        h, w, _ = frame_images[0].shape
+        out_path = "output/result.mp4"
+        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+        for img in frame_images:
+            writer.write(img)
+        writer.release()
+        return out_path
+
+    else:
+        raise ValueError("Unsupported file type")