import gradio as gr
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
import tempfile
import time
import os

MODEL_ID = "cerspense/zeroscope_v2_576w"

# Increase HF download timeout
os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"

def load_pipeline(retries=3):
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    for attempt in range(retries):
        try:
            print(f"Loading pipeline (attempt {attempt + 1}/{retries})…")
            pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype)
            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
            if torch.cuda.is_available():
                pipe = pipe.to("cuda")
                pipe.unet.enable_forward_chunking(chunk_size=1, dim=1)
                pipe.enable_vae_slicing()
            else:
                pipe = pipe.to("cpu")
            return pipe
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            if attempt < retries - 1:
                print("Retrying in 10 seconds…")
                time.sleep(10)
            else:
                raise

print("Loading AI Video Generator pipeline…")
pipe = load_pipeline()
print("Pipeline ready ✓")

IS_GPU = torch.cuda.is_available()
HW_WARNING = "" if IS_GPU else (
    "⚠️ **Running on CPU — videos will take 30–90 minutes.** "
    "Go to Space **Settings → Hardware → T4 Small** to enable GPU."
)

def generate_video(prompt, negative_prompt, num_inference_steps, guidance_scale, num_frames, fps, seed):
    if not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    generator = torch.Generator().manual_seed(int(seed)) if seed >= 0 else None
    all_frames = []
    chunk_size = 24

    if num_frames > chunk_size:
        for i in range(num_frames // chunk_size):
            result = pipe(
                prompt=f"{prompt}, continuous smooth motion",
                negative_prompt=negative_prompt or None,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                num_frames=chunk_size,
                height=320, width=576,
                generator=generator,
            ).frames[0]
            all_frames.extend(result)
        remainder = num_frames % chunk_size
        if remainder > 0:
            result = pipe(
                prompt=f"{prompt}, continuous smooth motion",
                negative_prompt=negative_prompt or None,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                num_frames=remainder,
                height=320, width=576,
                generator=generator,
            ).frames[0]
            all_frames.extend(result)
    else:
        all_frames = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            num_frames=num_frames,
            height=320, width=576,
            generator=generator,
        ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        out_path = tmp.name
    export_to_video(all_frames, out_path, fps=fps)
    return out_path


EXAMPLES = [
    ["A majestic eagle soaring over snow-capped mountains at golden hour, cinematic", "", 30, 7.5, 24, 8, 42],
    ["A futuristic city at night with neon lights reflecting on wet streets, cyberpunk", "blurry, low quality", 30, 7.5, 24, 8, 7],
    ["A timelapse of a blooming flower in a sunlit meadow, macro photography", "", 25, 6.5, 16, 8, 123],
]

with gr.Blocks(title="AI Powered 1 Minute Video Generator") as demo:

    gr.Markdown("# 🤖 AI Powered 1 Minute Video Generator")
    gr.Markdown("Generate stunning AI videos up to 1 full minute · Powered by ZeroScope V2")

    if HW_WARNING:
        gr.Markdown(HW_WARNING)

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Prompt", placeholder="Describe your video in detail…", lines=3)
            negative_prompt = gr.Textbox(label="Negative Prompt (optional)", placeholder="blurry, low quality, distorted…", lines=2)

            with gr.Row():
                num_steps = gr.Slider(10, 50, value=30, step=1, label="Inference Steps")
                guidance = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance Scale")

            with gr.Row():
                num_frames = gr.Slider(8, 480, value=24, step=1, label="Number of Frames")
                fps = gr.Slider(4, 16, value=8, step=1, label="FPS")

            seed = gr.Number(value=42, label="Seed (-1 = random)", precision=0)

            with gr.Row():
                ultra_btn   = gr.Button("🚀 Ultra Fast")
                fast_btn    = gr.Button("⚡ Fast Mode")
                quality_btn = gr.Button("🎨 Quality Mode")
                one_min_btn = gr.Button("🕐 1 Min Video")

            generate_btn = gr.Button("🎬 Generate Video", variant="primary")

        with gr.Column(scale=3):
            output_video = gr.Video(label="Generated Video", height=400)
            gr.Markdown("""
| Mode | Steps | Frames | FPS | Duration | T4 GPU |
|------|-------|--------|-----|----------|--------|
| 🚀 Ultra Fast | 10 | 8 | 8 | ~1 sec | ~30 sec |
| ⚡ Fast | 15 | 16 | 8 | ~2 sec | ~90 sec |
| 🎨 Quality | 30 | 24 | 8 | ~3 sec | ~3 min |
| 🕐 1 Min Video | 20 | 480 | 8 | ~60 sec | ~2 hrs |
            """)

    ultra_btn.click(fn=lambda: (10, 7.5, 8, 8),     inputs=[], outputs=[num_steps, guidance, num_frames, fps])
    fast_btn.click(fn=lambda: (15, 7.5, 16, 8),     inputs=[], outputs=[num_steps, guidance, num_frames, fps])
    quality_btn.click(fn=lambda: (30, 7.5, 24, 8),  inputs=[], outputs=[num_steps, guidance, num_frames, fps])
    one_min_btn.click(fn=lambda: (20, 7.5, 480, 8), inputs=[], outputs=[num_steps, guidance, num_frames, fps])

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt, negative_prompt, num_steps, guidance, num_frames, fps, seed],
        outputs=output_video,
        fn=generate_video,
        cache_examples=False,
    )

    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, negative_prompt, num_steps, guidance, num_frames, fps, seed],
        outputs=output_video,
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)