Update app.py

app.py

@@ -1,63 +1,101 @@
-import gradio as gr
+import os
 import cv2
+import gradio as gr
+import AnimeGANv3_src
 import numpy as np
-from diffusers import AutoPipelineForImage2Image
-from diffusers.utils import load_image
-from PIL import Image  # Add this import
- 
-# Load the anime-style diffusion model
-pipe = AutoPipelineForImage2Image.from_pretrained(
-    "nitrosocke/Arcane-Diffusion",
-    safety_checker=None,
-)
-# Running on CPU by default (no .to("cuda"))
-
-# Function to process a single frame
-def process_frame(frame, prompt):
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
-    pil_image = Image.fromarray(frame_rgb)  # Convert NumPy array to PIL image
-    image = load_image(pil_image)  # Pass PIL image to load_image
-    result = pipe(prompt=prompt, image=image, strength=0.75).images[0]
-    return np.array(result)
-
-# Function to convert the entire video
-def video_to_anime(video_path, prompt="Arcane style"):
-    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
-    processed_frames = [process_frame(frame, prompt) for frame in frames]
-
-    # Write the output video
-    height, width, _ = processed_frames[0].shape
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_path = "output.mp4"
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-    for frame in processed_frames:
-        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-        out.write(frame_bgr)
-    out.release()
-
-    return output_path
-
-# Create the Gradio interface
+from moviepy.editor import VideoFileClip, ImageSequenceClip
+
+class AnimeGANv3:
+    def __init__(self):
+        # Ensure the output directory exists
+        os.makedirs('output', exist_ok=True)
+
+    def process_frame(self, frame, style_code, det_face):
+        """Process a single frame with AnimeGANv3."""
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        output = AnimeGANv3_src.Convert(frame_rgb, style_code, det_face)
+        return output
+
+    def inference(self, video_path, style, if_face=None):
+        print(video_path, style, if_face)
+        try:
+            # Map style names to codes
+            style_codes = {
+                "AnimeGANv3_Arcane": "A",
+                "AnimeGANv3_Trump v1.0": "T",
+                "AnimeGANv3_Shinkai": "S",
+                "AnimeGANv3_PortraitSketch": "P",
+                "AnimeGANv3_Hayao": "H",
+                "AnimeGANv3_Disney v1.0": "D",
+                "AnimeGANv3_JP_face v1.0": "J",
+                "AnimeGANv3_Kpop v2.0": "K",
+            }
+            style_code = style_codes.get(style, "U")
+            det_face = if_face == "Yes"
+
+            # Open the video
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                raise Exception("Could not open video file")
+
+            # Get video properties
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            print(f"Processing {frame_count} frames at {fps} FPS")
+
+            # Extract and process frames
+            frames = []
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                processed_frame = self.process_frame(frame, style_code, det_face)
+                frames.append(processed_frame)
+            cap.release()
+
+            # Convert frames to video
+            save_path = "output/out.mp4"
+            clip = ImageSequenceClip(frames, fps=fps)
+            clip.write_videofile(save_path, codec="libx264", audio=False)
+
+            return save_path
+        except Exception as error:
+            print('Error:', error)
+            return None
+
+# Create an instance of the AnimeGANv3 class
+anime_gan = AnimeGANv3()
+
+# Define the Gradio interface
+title = "AnimeGANv3: Video to Anime Converter"
+description = r"""Upload a video to convert it into anime style using AnimeGANv3.<br>
+Select a style and choose whether to optimize for faces.<br>
+<a href='https://github.com/TachibanaYoshino/AnimeGANv3' target='_blank'><b>AnimeGANv3 GitHub</b></a> | 
+<a href='https://www.patreon.com/Asher_Chan' target='_blank'><b>Patreon</b></a>"""
+
 iface = gr.Interface(
-    fn=video_to_anime,
+    fn=anime_gan.inference,
     inputs=[
         gr.Video(label="Input Video"),
-        gr.Textbox(label="Style Prompt", value="Arcane style")
+        gr.Dropdown(choices=[
+            'AnimeGANv3_Hayao',
+            'AnimeGANv3_Shinkai',
+            'AnimeGANv3_Arcane',
+            'AnimeGANv3_Trump v1.0',
+            'AnimeGANv3_Disney v1.0',
+            'AnimeGANv3_PortraitSketch',
+            'AnimeGANv3_JP_face v1.0',
+            'AnimeGANv3_Kpop v2.0',
+        ], label='AnimeGANv3 Style', default='AnimeGANv3_Arcane'),
+        gr.Radio(choices=["Yes", "No"], label='Extract face', default="No"),
+    ],
+    outputs=[
+        gr.Video(label="Output Video")
     ],
-    outputs=gr.Video(label="Output Video"),
-    title="Video to Anime Converter",
-    description="Upload a video and convert it to anime style!"
+    title=title,
+    description=description,
+    allow_flagging="never"
 )
 
-# Launch the interface with a public link
-iface.launch(share=True)  # Added share=True as per the suggestion
+# Launch the interface
+iface.launch()