coco / app.py
Mavhas's picture
Update app.py
d1af892 verified
raw
history blame contribute delete
726 Bytes
import imageio
import imageio_ffmpeg
import torch
from diffusers import MochiPipeline
from diffusers.utils import export_to_video
# Load the pre-trained video generation model
model = MochiPipeline.from_pretrained(
"MISHANM/video_generation",
# variant="bf16",
torch_dtype=torch.bfloat16,
device_map="balanced"
)
# Enable memory savings by tiling the VAE
model.enable_vae_tiling()
# Define the prompt and number of frames
prompt = "A cow drinking water on the surface of Mars."
num_frames = 20
frames = model(prompt, num_frames=num_frames).frames[0]
export_to_video(frames, "video.mp4", fps=30)
print("Video generation complete. Saved as 'video.mp4'.")