test / app.py
reshav1's picture
Update app.py
72e95a8 verified
raw
history blame
1.72 kB
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
import streamlit as st
from diffusers import UNet2DConditionModel, TextEncoder, VQModel
# Use the default model names here
unet_model_name = "unet/diffusion_pytorch_model.bin"
text_encoder_name = "text_encoder/pytorch_model.bin"
vae_model_name = "vae/diffusion_pytorch_model.bin"
# Create the pipeline or model objects using the default names
pipeline = UNet2DConditionModel.from_pretrained(unet_model_name)
# Title and User Input
st.title("Text-to-Video with Streamlit")
prompt = st.text_input("Enter your text prompt:", "Spiderman is surfing")
# Button to trigger generation
if st.button("Generate Video"):
# Ensure you have 'accelerate' version 0.17.0 or higher (see previous explanation)
import accelerate
if accelerate.__version__ < "0.17.0":
st.warning("Please upgrade 'accelerate' to version 0.17.0 or higher for CPU offloading.")
else:
with st.spinner("Generating video..."):
pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w",
torch_dtype=torch.float16,
variant="fp16",
device="cpu") # Force CPU usage
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload() # Assuming 'accelerate' is updated
video_frames = pipe(prompt, num_inference_steps=25).frames
video_path = export_to_video(video_frames)
# Display the video in the Streamlit app
st.video(video_path)