import streamlit as st
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video

# Explicitly set the device to CPU
device = torch.device("cpu")

# Load the model onto the CPU
pipe = DiffusionPipeline.from_pretrained(
    "damo-vilab/text-to-video-ms-1.7b", 
    torch_dtype=torch.float32  # Use float32 for CPU
).to(device)

# Initialize the scheduler from the pipeline's configuration, no need to move it to the CPU
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

prompt = "Pop international experimental music"

# Generate the video frames on the CPU
video_frames = pipe(prompt, num_inference_steps=25).frames

# Export the frames to a video file
video_path = export_to_video(video_frames)