import gradio as gr
import spaces
import torch
import tempfile

from diffusers import LTXImageToVideoPipeline
from diffusers.utils import export_to_video

pipe = LTXImageToVideoPipeline.from_pretrained(
    "Lightricks/LTX-Video",
    torch_dtype=torch.bfloat16
)

pipe.to("cuda")


@spaces.GPU
def generate_video(image, prompt):
    if image is None:
        raise gr.Error("Please upload an image.")

    if not prompt.strip():
        raise gr.Error("Please describe the movement.")

    video = pipe(
        image=image,
        prompt=prompt,
        width=704,
        height=480,
        num_frames=49,
        num_inference_steps=20,
    ).frames[0]

    output_file = tempfile.NamedTemporaryFile(
        suffix=".mp4",
        delete=False
    ).name

    export_to_video(video, output_file, fps=24)

    return output_file


demo = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Image(type="pil", label="Real Image"),
        gr.Textbox(
            label="Motion Prompt",
            placeholder="Describe realistic movement..."
        ),
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Real Image → Video",
)

demo.launch()