import gradio as gr
import torch
from diffusers import I2VGenXLPipeline
from diffusers.utils import export_to_gif, load_image

def generate_video(prompt, image_url, negative_prompt, num_inference_steps, guidance_scale, seed):
    try:
        # Load pipeline
        pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
        pipeline.enable_model_cpu_offload()

        # Load image
        image = load_image(image_url).convert("RGB")

        # Set random seed
        generator = None if seed == -1 else torch.manual_seed(seed)

        # Generate frames
        result = pipeline(
            prompt=prompt,
            image=image,
            num_inference_steps=num_inference_steps,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            generator=generator,
        ).frames[0]

        # Export to GIF
        gif_path = "i2v_output.gif"
        export_to_gif(result, gif_path)

        return gif_path
    except Exception as e:
        return f"An error occurred: {e}"

# Gradio Interface
interface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Enter your prompt", lines=2, placeholder="Describe your desired video."),
        gr.Textbox(label="Enter image URL", placeholder="Paste the URL of the input image."),
        gr.Textbox(label="Enter negative prompt (optional)", placeholder="Describe undesired aspects (e.g., blurry, distorted)."),
        gr.Slider(label="Number of Inference Steps", minimum=10, maximum=100, value=50, step=1),
        gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.1),
        gr.Number(label="Seed (set to -1 for random)", value=8888),
    ],
    outputs=gr.Image(label="Generated Video (GIF)", type="filepath"),
    title="Image-to-Video Generator",
    description="Generate videos from text prompts and input images using AI.",
    article="Made by: Mostafa Hazem."
)

if __name__ == "__main__":
    interface.launch()