try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(duration=None):
            def decorator(func):
                return func
            return decorator

import gradio as gr
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
import torch
import random

# Load pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to(device)

# Styles map
STYLES = {
    "None": "{prompt}",
    "Cinematic": "{prompt}, cinematic style, highly detailed, photorealistic, 8k resolution, dramatic volumetric lighting, depth of field",
    "3D Animation": "{prompt}, 3D Pixar style character, vibrant colors, clean textures, whimsical, ray-traced shadows",
    "Cyberpunk": "{prompt}, cyberpunk aesthetic, glowing neon lights, rain-slicked streets, futuristic atmosphere, high contrast",
    "Anime": "{prompt}, modern anime style, beautiful hand-drawn aesthetics, soft color grading, high detail, studio Ghibli influence"
}

@spaces.GPU(duration=120)  # seconds of GPU time this function may use
def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, resolution, num_frames, style, seed, randomize_seed):
    # Apply style template
    styled_prompt = STYLES.get(style, "{prompt}").format(prompt=prompt)
    
    # Parse resolution (e.g. "768x512")
    width, height = map(int, resolution.split("x"))
    
    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
        
    generator = torch.Generator(device="cpu").manual_seed(seed)
    
    video = pipe(
        prompt=styled_prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_frames=int(num_frames),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        generator=generator
    ).frames[0]
    
    export_to_video(video, "output.mp4", fps=24)
    return "output.mp4", seed

# Custom Gradio UI with advanced quality controls
with gr.Blocks(title="LTX-Video Generator Pro") as demo:
    gr.Markdown("# 🎬 LTX-Video Text-to-Video Generator")
    gr.Markdown("Generate high-quality videos using Lightricks LTX-Video on Hugging Face ZeroGPU.")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="A cinematic shot of a sunset over the ocean, high quality, 4k",
                lines=3
            )
            style = gr.Dropdown(
                label="Prompt Style Preset",
                choices=list(STYLES.keys()),
                value="None"
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Aids Quality)", 
                value="worst quality, low quality, deformed, distorted, blurry, noisy, static, cartoon, lowres",
                lines=2
            )
            
            with gr.Accordion("Advanced Settings (Quality Controls)", open=True):
                resolution = gr.Dropdown(
                    label="Resolution", 
                    choices=["768x512", "512x768", "768x768", "960x544"], 
                    value="768x512"
                )
                num_frames = gr.Slider(
                    label="Number of Frames (Multiple of 8 + 1)", 
                    minimum=17, 
                    maximum=121, 
                    step=8, 
                    value=65
                )
                num_inference_steps = gr.Slider(
                    label="Inference Steps (Higher = more detail)", 
                    minimum=10, 
                    maximum=50, 
                    step=1, 
                    value=30
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale (Prompt adherence)", 
                    minimum=1.0, 
                    maximum=10.0, 
                    step=0.5, 
                    value=3.0
                )
                seed = gr.Number(
                    label="Seed", 
                    value=42, 
                    precision=0
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize Seed on Generate", 
                    value=True
                )
            
            generate_btn = gr.Button("Generate Video", variant="primary")
            
        with gr.Column(scale=1):
            output_video = gr.Video(label="Generated Video")
            output_seed = gr.Number(label="Used Seed")
            
    generate_btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, resolution, num_frames, style, seed, randomize_seed],
        outputs=[output_video, output_seed]
    )

demo.launch()