import streamlit as st
import torch
from diffusers import DiffusionPipeline
import tempfile

# Load the text-to-video model
st.write("Loading model... (first run may take a few minutes)")

model_id = "damo-vilab/text-to-video-ms-1.7b"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cpu")  # Stay on CPU since we don’t have a GPU

st.title("Text-to-Video Generator")
prompt = st.text_input("Enter a text prompt for the video:")

frames = st.slider("Number of frames (video length)", min_value=8, max_value=24, value=16)

if st.button("Generate Video") and prompt:
    with st.spinner("Generating video... this may take a while on CPU"):
        result = pipe(prompt=prompt, num_frames=frames, num_inference_steps=20)
        video_frames = result.frames  # List of PIL images

        # Save frames as video file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
            video_path = temp_file.name
        result.export_to_video(video_path, fps=8)

    st.video(video_path)