import spaces  # <--- YEH LINE SAB SE UPAR HONGII REQUIRED FOR ZEROGPU
import os
import torch
import gradio as gr
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

MODEL_ID = "FreeAi1111/PrepGenius"
HF_TOKEN = os.getenv("HF_TOKEN")

print("Loading PrepGenius CogVideoX Model...")
pipe = CogVideoXPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN
)
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()

@spaces.GPU(duration=120)
def generate_video(prompt, duration_choice):
    frames = 49 if duration_choice == "6 Seconds" else 81
    
    video = pipe(
        prompt=prompt,
        num_videos_per_prompt=1,
        num_inference_steps=50,
        num_frames=frames,
        guidance_scale=6.0,
        generator=torch.Generator(device="cuda").manual_seed(42)
    ).frames[0]

    output_path = "output.mp4"
    export_to_video(video, output_path, fps=8)
    return output_path

# UI Interface
with gr.Blocks(title="PrepGenius AI Video Studio") as demo:
    gr.Markdown("# 🎬 PrepGenius AI Video Generator")
    gr.Markdown("Powered by `FreeAi1111/PrepGenius` (CogVideoX-5B)")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Video Prompt", 
                placeholder="A cinematic drone shot of a futuristic cyberpunk city at sunset...", 
                lines=3
            )
            duration_input = gr.Radio(
                ["6 Seconds", "10 Seconds"], 
                label="Duration", 
                value="6 Seconds"
            )
            generate_btn = gr.Button("Generate Video", variant="primary")
        
        with gr.Column():
            video_output = gr.Video(label="Generated Result")

    generate_btn.click(
        fn=generate_video,
        inputs=[prompt_input, duration_input],
        outputs=video_output
    )

demo.queue().launch()