import torch
from diffusers import DiffusionPipeline
import gradio as gr
import os

# 1. Load the Model
# Use 'cuda' if you have an NVIDIA GPU and CUDA installed, otherwise 'cpu'
# Hugging Face Spaces will handle the GPU for you if you selected it
device = "cuda" if torch.cuda.is_available() else "cpu"

try:
    pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
    pipe.to(device)
    print(f"Model loaded successfully on {device}!")
except Exception as e:
    print(f"Error loading model: {e}")
    print("Ensure all dependencies are in requirements.txt and GPU is available if using 'cuda'.")
    pipe = None # Set pipe to None if loading fails

# 2. Define the Generation Function
def generate_video(prompt: str, num_frames: int = 24, guidance_scale: float = 9.0):
    if pipe is None:
        return "Error: AI model not loaded. Please check server logs in the 'Logs' tab."

    if not prompt:
        return "Please enter a text prompt."

    print(f"Generating video for prompt: '{prompt}'")
    try:
        # Generate video frames
        video_frames = pipe(
            prompt=prompt,
            num_frames=num_frames,
            guidance_scale=guidance_scale
        ).frames

        # The output file path for the video
        output_filename = "generated_video.mp4"
        # pipe.to_video creates the video file from frames
        video_path = pipe.to_video(video_frames, output_filename)
        return video_path
    except Exception as e:
        print(f"Video generation failed: {e}")
        return f"Video generation failed: {e}"

# 3. Create Gradio Interface
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Text Prompt", placeholder="e.g., 'A robot dancing in the rain'"),
        gr.Slider(minimum=16, maximum=64, step=4, value=24, label="Number of Frames (Video Length)"),
        gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=9.0, label="Guidance Scale (Creativity vs. Prompt Adherence)")
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Simple Text-to-Video Generator (Zeroscope v2)",
    description="Enter a text prompt to generate a short video using Zeroscope v2."
)

# 4. Launch the App
if __name__ == "__main__":
    iface.launch() # Hugging Face Spaces handles sharing automatically