import streamlit as st
from diffusers import DiffusionPipeline
import torch
from moviepy.editor import *
import numpy as np
import tempfile, os

st.title("🚀 Text-to-Video (Zeroscope)")

@st.cache_resource
def load_model():
    pipe = DiffusionPipeline.from_pretrained(
        "cerspense/zeroscope_v2_576w", 
        torch_dtype=torch.float16
    )
    pipe.enable_cpu_offload()
    return pipe

pipe = load_model()

prompt = st.text_area("Enter prompt (short & descriptive):", max_chars=50)

if st.button("Generate Video"):
    if prompt:
        with st.spinner("Generating... (this may take ~2-3 mins)"):
            video_frames = pipe(prompt, num_frames=10, height=320, width=576).frames

            video_filename = tempfile.mktemp(".mp4")
            clips = [ImageClip(np.array(frame)).set_duration(0.3) for frame in video_frames]
            final_clip = concatenate_videoclips(clips, method="compose")
            final_clip.write_videofile(video_filename, fps=5)

            st.video(video_filename)

            os.remove(video_filename)
    else:
        st.warning("Enter a prompt to generate a video.")