import gradio as gr
import cv2
import numpy as np
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

# Load the anime-style diffusion model
pipe = AutoPipelineForImage2Image.from_pretrained(
    "nitrosocke/Arcane-Diffusion",
    safety_checker=None,
)
pipe.to("cuda")

# Function to process a single frame
def process_frame(frame, prompt):
    # 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"):
    # 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 to anime style
    processed_frames = [process_frame(frame, prompt) for frame in frames]

    # Write the output video using OpenCV
    height, width, _ = processed_frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
    output_path = "output.mp4"
    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

# Create the Gradio interface
iface = gr.Interface(
    fn=video_to_anime,
    inputs=[
        gr.Video(label="Input Video"),
        gr.Textbox(label="Style Prompt", default="Arcane style")
    ],
    outputs=gr.Video(label="Output Video"),
    title="Video to Anime Converter",
    description="Upload a video and convert it to anime style!"
)

# Launch the interface
iface.launch()