Skip to content

Commit 0a2c26d

Browse files
Update Documentation for NVIDIA Cosmos (#13251)
* fix docs * update main example
1 parent 07c5ba8 commit 0a2c26d

File tree

4 files changed

+24
-18
lines changed

4 files changed

+24
-18
lines changed

docs/source/en/_toctree.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -532,8 +532,6 @@
532532
title: ControlNet-XS with Stable Diffusion XL
533533
- local: api/pipelines/controlnet_union
534534
title: ControlNetUnion
535-
- local: api/pipelines/cosmos
536-
title: Cosmos
537535
- local: api/pipelines/ddim
538536
title: DDIM
539537
- local: api/pipelines/ddpm
@@ -677,6 +675,8 @@
677675
title: CogVideoX
678676
- local: api/pipelines/consisid
679677
title: ConsisID
678+
- local: api/pipelines/cosmos
679+
title: Cosmos
680680
- local: api/pipelines/framepack
681681
title: Framepack
682682
- local: api/pipelines/helios

docs/source/en/api/pipelines/cosmos.md

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,29 +21,31 @@
2121
> [!TIP]
2222
> Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines.
2323
24-
## Loading original format checkpoints
25-
26-
Original format checkpoints that have not been converted to diffusers-expected format can be loaded using the `from_single_file` method.
24+
## Basic usage
2725

2826
```python
2927
import torch
30-
from diffusers import Cosmos2TextToImagePipeline, CosmosTransformer3DModel
31-
32-
model_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
33-
transformer = CosmosTransformer3DModel.from_single_file(
34-
"https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image/blob/main/model.pt",
35-
torch_dtype=torch.bfloat16,
36-
).to("cuda")
37-
pipe = Cosmos2TextToImagePipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)
28+
from diffusers import Cosmos2_5_PredictBasePipeline
29+
from diffusers.utils import export_to_video
30+
31+
model_id = "nvidia/Cosmos-Predict2.5-2B"
32+
pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
33+
model_id, revision="diffusers/base/post-trained", torch_dtype=torch.bfloat16
34+
)
3835
pipe.to("cuda")
3936

40-
prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
37+
prompt = "As the red light shifts to green, the red bus at the intersection begins to move forward, its headlights cutting through the falling snow. The snowy tire tracks deepen as the vehicle inches ahead, casting fresh lines onto the slushy road. Around it, streetlights glow warmer, illuminating the drifting flakes and wet reflections on the asphalt. Other cars behind start to edge forward, their beams joining the scene. The stillness of the urban street transitions into motion as the quiet snowfall is punctuated by the slow advance of traffic through the frosty city corridor."
4138
negative_prompt = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
4239

4340
output = pipe(
44-
prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator().manual_seed(1)
45-
).images[0]
46-
output.save("output.png")
41+
image=None,
42+
video=None,
43+
prompt=prompt,
44+
negative_prompt=negative_prompt,
45+
num_frames=93,
46+
generator=torch.Generator().manual_seed(1),
47+
).frames[0]
48+
export_to_video(output, "text2world.mp4", fps=16)
4749
```
4850

4951
## Cosmos2_5_TransferPipeline

docs/source/en/api/pipelines/overview.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ The table below lists all the pipelines currently available in 🤗 Diffusers an
4444
| [ControlNet with Stable Diffusion XL](controlnet_sdxl) | text2image |
4545
| [ControlNet-XS](controlnetxs) | text2image |
4646
| [ControlNet-XS with Stable Diffusion XL](controlnetxs_sdxl) | text2image |
47+
| [Cosmos](cosmos) | text2video, video2video |
4748
| [Dance Diffusion](dance_diffusion) | unconditional audio generation |
4849
| [DDIM](ddim) | unconditional image generation |
4950
| [DDPM](ddpm) | unconditional image generation |

src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,13 +82,16 @@ def retrieve_latents(
8282
```python
8383
>>> import cv2
8484
>>> import numpy as np
85+
>>> from PIL import Image
8586
>>> import torch
8687
>>> from diffusers import Cosmos2_5_TransferPipeline, AutoModel
8788
>>> from diffusers.utils import export_to_video, load_video
8889
8990
>>> model_id = "nvidia/Cosmos-Transfer2.5-2B"
9091
>>> # Load a Transfer2.5 controlnet variant (edge, depth, seg, or blur)
91-
>>> controlnet = AutoModel.from_pretrained(model_id, revision="diffusers/controlnet/general/edge")
92+
>>> controlnet = AutoModel.from_pretrained(
93+
... model_id, revision="diffusers/controlnet/general/edge", torch_dtype=torch.bfloat16
94+
... )
9295
>>> pipe = Cosmos2_5_TransferPipeline.from_pretrained(
9396
... model_id, controlnet=controlnet, revision="diffusers/general", torch_dtype=torch.bfloat16
9497
... )

0 commit comments

Comments
 (0)