Update app.py

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 import cv2
-from moviepy.editor import VideoFileClip, ImageSequenceClip
 import numpy as np
 from diffusers import AutoPipelineForImage2Image
 from diffusers.utils import load_image
 
-# Load the anime-style model
+# Load the anime-style diffusion model
 pipe = AutoPipelineForImage2Image.from_pretrained(
     "nitrosocke/Arcane-Diffusion",
     safety_checker=None,
@@ -14,27 +13,41 @@ pipe.to("cuda")
 
 # Function to process a single frame
 def process_frame(frame, prompt):
-    # Convert frame from BGR (OpenCV) to RGB
-    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    # Load the frame as an image for the model
-    image = load_image(frame)
+    # Convert frame from BGR (OpenCV format) to RGB (expected by the model)
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    # Load the frame as an image for the diffusion model
+    image = load_image(frame_rgb)
     # Apply the anime-style transformation
     result = pipe(prompt=prompt, image=image, strength=0.75).images[0]
+    # Convert back to numpy array
     return np.array(result)
 
 # Function to convert the entire video
 def video_to_anime(video_path, prompt="Arcane style"):
-    # Load the video and extract frames
-    clip = VideoFileClip(video_path)
-    frames = [frame for frame in clip.iter_frames()]
+    # Read the input video using OpenCV
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(frame)
+    cap.release()
 
-    # Process each frame with the anime-style model
+    # Process each frame to anime style
     processed_frames = [process_frame(frame, prompt) for frame in frames]
 
-    # Reassemble the processed frames into a video
-    new_clip = ImageSequenceClip(processed_frames, fps=clip.fps)
+    # Write the output video using OpenCV
+    height, width, _ = processed_frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
     output_path = "output.mp4"
-    new_clip.write_videofile(output_path, codec="libx264")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    for frame in processed_frames:
+        # Convert back to BGR for OpenCV
+        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        out.write(frame_bgr)
+    out.release()
 
     return output_path