import streamlit as st
from diffusers import DiffusionPipeline
import torch
import os
import imageio

st.set_page_config(page_title="Gen-V", layout="centered")
st.title("🎥 Gen-V: Text-to-Video Generator")
st.write("Generate AI-powered videos from text prompts using open-source models!")

prompt = st.text_input("Enter a prompt", "cat wearing black goggles")

if st.button("🎬 Generate Video"):
    with st.spinner("Generating video... this might take a bit ⏳"):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        use_fp16 = torch.cuda.is_available()

        pipe = DiffusionPipeline.from_pretrained(
            "damo-vilab/text-to-video-ms-1.7b",
            torch_dtype=torch.float16 if use_fp16 else torch.float32,
            variant="fp16" if use_fp16 else None,
        ).to("cuda" if use_fp16 else "cpu")

        video_frames = pipe(prompt).frames  # Returns list of PIL Images (frames)

        output_path = "genv_output.mp4"
        imageio.mimsave(output_path, video_frames, fps=8)

        st.success("Video generation complete!")
        st.video(output_path)