|
|
|
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") |
