import torch
import streamlit as st
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

st.title("CogVideoX Video Generation on CPU")

# Your prompt
prompt = "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance."

if st.button("Generate Video"):
    with st.spinner("Generating video on CPU..."):
        # Load the model and force it to run on CPU
        pipe = CogVideoXPipeline.from_pretrained(
            "THUDM/CogVideoX-2b",
            torch_dtype=torch.float32  # Use float32 for CPU
        )

        # Ensure model uses CPU
        pipe.to("cpu")

        # VAE optimizations
        pipe.vae.enable_slicing()
        pipe.vae.enable_tiling()

        # Generate the video
        video = pipe(
            prompt=prompt,
            num_videos_per_prompt=1,
            num_inference_steps=50,
            num_frames=49,
            guidance_scale=6,
            generator=torch.manual_seed(42),  # Use manual_seed for CPU
        ).frames[0]

        # Export video
        export_to_video(video, "output.mp4", fps=8)
        
        # Show the video in Streamlit
        st.video("output.mp4")

        st.success("Video generated successfully on CPU!")