Skip to content

Commit 6025848

Browse files
authored
Update pipeline_stg_mochi.py
1 parent 630f116 commit 6025848

File tree

1 file changed

+15
-4
lines changed

1 file changed

+15
-4
lines changed

examples/community/pipeline_stg_mochi.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,24 @@
5656
>>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16)
5757
>>> pipe.enable_model_cpu_offload()
5858
>>> pipe.enable_vae_tiling()
59-
>>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
60-
>>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0]
59+
>>> prompt = "A close-up of a beautiful woman's face with colored powder exploding around her, creating an abstract splash of vibrant hues, realistic style."
60+
61+
>>> # Configure STG mode options
62+
>>> stg_applied_layers_idx = [34] # Layer indices from 0 to 41
63+
>>> stg_scale = 1.0
64+
>>> do_rescaling = False
65+
66+
>>> frames = pipe(prompt,
67+
... prompt=prompt,
68+
... num_inference_steps=28,
69+
... guidance_scale=3.5,
70+
... stg_applied_layers_idx=stg_applied_layers_idx,
71+
... stg_scale=stg_scale,
72+
... do_rescaling=do_rescaling).frames[0]
6173
>>> export_to_video(frames, "mochi.mp4")
6274
```
6375
"""
6476

65-
6677
def forward_with_stg(
6778
self,
6879
hidden_states: torch.Tensor,
@@ -822,4 +833,4 @@ def __call__(
822833
if not return_dict:
823834
return (video,)
824835

825-
return MochiPipelineOutput(frames=video)
836+
return MochiPipelineOutput(frames=video)

0 commit comments

Comments
 (0)