| | |
| | from textwrap import dedent |
| | from genmo.mochi_preview.pipelines import ( |
| | DecoderModelFactory, |
| | DitModelFactory, |
| | MochiSingleGPUPipeline, |
| | T5ModelFactory, |
| | linear_quadratic_schedule, |
| | ) |
| | from genmo.lib.utils import save_video |
| | from genmo.lib.progress import progress_bar |
| | from pathlib import Path |
| | import sys |
| |
|
| | MOCHI_DIR = sys.argv[1] |
| | assert Path(MOCHI_DIR).exists(), f"Model directory {MOCHI_DIR} does not exist." |
| | pipeline = MochiSingleGPUPipeline( |
| | text_encoder_factory=T5ModelFactory(), |
| | dit_factory=DitModelFactory(model_path=f"{MOCHI_DIR}/dit.safetensors", model_dtype="bf16"), |
| | decoder_factory=DecoderModelFactory( |
| | model_path=f"{MOCHI_DIR}/vae.safetensors", |
| | model_stats_path=f"{MOCHI_DIR}/vae_stats.json", |
| | ), |
| | cpu_offload=True, |
| | decode_type="tiled_full" |
| | ) |
| |
|
| | PROMPT = dedent(""" |
| | A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl |
| | filled with lemons and sprigs of mint against a peach-colored background. |
| | The hand gently tosses the lemon up and catches it, showcasing its smooth texture. |
| | A beige string bag sits beside the bowl, adding a rustic touch to the scene. |
| | Additional lemons, one halved, are scattered around the base of the bowl. |
| | The even lighting enhances the vibrant colors and creates a fresh, |
| | inviting atmosphere. |
| | """) |
| |
|
| | video = pipeline( |
| | height=480, |
| | width=848, |
| | num_frames=31, |
| | num_inference_steps=64, |
| | sigma_schedule=linear_quadratic_schedule(64, 0.025), |
| | cfg_schedule=[4.5] * 64, |
| | batch_cfg=False, |
| | prompt=PROMPT, |
| | negative_prompt="", |
| | seed=12345, |
| | ) |
| |
|
| | with progress_bar(type="tqdm"): |
| | save_video(video[0], "video.mp4") |
