Cosmos3 ModularPipeline (#14110)

* Cosmos3 ModularPipeline initial commit

* Fix from_pretrained for modular pipleine without modular_model_index.json

* Isolate text prep inside Cosmos3TextEncoderStep; add Cosmos3ActionTextStep

* Cosmos3 modular encode/decode flow

* Drop device/dtype required InputParams; use components._execution_device and model dtype

* Add lazy pipeline properties and requires_safety_checker-based safety gating

* Drop callback_on_step_end from Cosmos3 denoise loop

* Move container/auto blocks into modular_blocks

* Move action_gen check, remove unused params

* Declare video_processor as ComponentSpec; drop unused pipeline state

* Split VAE encode block by modality (image and video)

* Remove Cosmos3AssembleResultStep; add Cosmos3ExportStep

* Remove unused action_mode

* Refactor Cosmos3 modular denoising by modality

* Document Cosmos3 modular output selection

* Address Cosmos3 modular review feedback

* Remaining local-variable device cleanup

* Address Cosmos3 modular review feedback

* Address final Cosmos3 modular feedback

---------

Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>

Author: yzhautouskay

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

@@ -738,10 +738,211 @@ pipe = Cosmos3OmniPipeline.from_pretrained(
 - all
 - __call__
 
+## Cosmos3OmniModularPipeline
+
+Cosmos 3 is also available as a Modular Diffusers pipeline. The task-based [`Cosmos3OmniPipeline`] remains available; the modular pipeline coexists with it and covers the same modes (`text2image`, `text2video`, `image2video`, `video2video`, and action-conditioned generation, with optional sound when supported by the checkpoint).
+
+```python
+import torch
+from diffusers import Cosmos3OmniModularPipeline
+
+pipe = Cosmos3OmniModularPipeline.from_pretrained(
+    "nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16
+)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.enable_safety_checker()
+
+videos = pipe(
+    prompt='{"scene":"A robot arm in a kitchen"}',
+    num_frames=1,
+    height=720,
+    width=1280,
+    output="videos",
+)
+
+# Modular pipelines expose declared outputs directly instead of using the task pipeline's
+# `return_dict`/`Cosmos3OmniPipelineOutput` API.
+image = videos[0]
+```
+
+You can also load through [`ModularPipeline`] and let the repository config select the blocks class:
+
+```python
+import torch
+from diffusers import ModularPipeline
+
+pipe = ModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.enable_safety_checker()
+videos = pipe(
+    prompt='{"scene":"A robot arm in a kitchen"}', num_frames=1, height=720, width=1280, output="videos"
+)
+```
+
+To inspect or customize a specific Cosmos modular workflow, use `available_workflows` + `get_workflow()`:
+
+```python
+available = pipe.blocks.available_workflows
+image2video_blocks = pipe.blocks.get_workflow("image2video")
+```
+
+### Modular examples for all existing workflows
+
+The modular pipeline supports the same call signatures as the task pipeline. The snippets below mirror every generation example shown above (`text2video`, `text2image`, `image2video`, `video2video`, `video2video_sound`, `text2video_sound`, and `action_policy`).
+
+```python
+import json
+import torch
+from diffusers import Cosmos3OmniModularPipeline, CosmosActionCondition
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import encode_video, export_to_video, load_image, load_video
+
+pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.enable_safety_checker()
+pipe.to("cuda")
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
+)
+
+# text2video
+json_prompt = json.load(open("assets/example_t2v_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt.json"))
+videos = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    num_frames=189,
+    height=720,
+    width=1280,
+    num_inference_steps=35,
+    guidance_scale=6.0,
+    fps=24.0,
+    output="videos",
+)
+export_to_video(videos, "cosmos3_modular_t2v.mp4", fps=24, macro_block_size=1)
+
+# text2image
+json_prompt = json.load(open("assets/example_t2i_prompt.json"))
+videos = pipe(prompt=json.dumps(json_prompt), num_frames=1, height=720, width=1280, output="videos")
+videos[0].save("cosmos3_modular_t2i.jpg", format="JPEG", quality=85)
+
+# image2video
+json_prompt = json.load(open("assets/example_i2v_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
+image = load_image("https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg")
+videos = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    image=image,
+    num_frames=189,
+    height=720,
+    width=1280,
+    fps=24.0,
+    output="videos",
+)
+export_to_video(videos, "cosmos3_modular_i2v.mp4", fps=24, macro_block_size=1)
+
+# video2video
+json_prompt = json.load(open("assets/example_v2v_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
+video = load_video(
+    "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
+)
+videos = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    video=video,
+    condition_frame_indexes_vision=[0, 1],
+    condition_video_keep="first",
+    num_frames=189,
+    height=720,
+    width=1280,
+    num_inference_steps=35,
+    guidance_scale=6.0,
+    fps=24.0,
+    output="videos",
+)
+export_to_video(videos, "cosmos3_modular_v2v.mp4", fps=24, macro_block_size=1)
+
+# video2video_sound
+outputs = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    video=video,
+    condition_frame_indexes_vision=[0, 1],
+    condition_video_keep="first",
+    num_frames=189,
+    height=720,
+    width=1280,
+    fps=24.0,
+    enable_sound=True,
+    output=["videos", "sound", "sampling_rate"],
+)
+encode_video(
+    outputs["videos"],
+    fps=24,
+    audio=outputs["sound"],
+    audio_sample_rate=outputs["sampling_rate"],
+    output_path="cosmos3_modular_v2v_with_sound.mp4",
+)
+
+# text2video_sound
+json_prompt = json.load(open("assets/example_t2v_sound_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt.json"))
+outputs = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    num_frames=189,
+    height=720,
+    width=1280,
+    fps=24.0,
+    enable_sound=True,
+    output=["videos", "sound", "sampling_rate"],
+)
+encode_video(
+    outputs["videos"],
+    fps=24,
+    audio=outputs["sound"],
+    audio_sample_rate=outputs["sampling_rate"],
+    output_path="cosmos3_modular_t2v_with_sound.mp4",
+)
+
+# action_policy
+prompt = "Put the pot to the left of the purple item."
+action_video = load_video(
+    "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_20260501_0.mp4"
+)
+outputs = pipe(
+    prompt=prompt,
+    action=CosmosActionCondition(
+        mode="policy",
+        chunk_size=16,
+        domain_name="bridge_orig_lerobot",
+        resolution_tier=480,
+        video=action_video,
+        view_point="ego_view",
+    ),
+    fps=5,
+    num_inference_steps=30,
+    guidance_scale=1.0,
+    use_system_prompt=False,
+    output=["videos", "action"],
+)
+export_to_video(outputs["videos"], "cosmos3_modular_action_policy.mp4", fps=5, macro_block_size=1)
+if outputs["action"] is not None:
+    with open("cosmos3_modular_action_policy.json", "w") as f:
+        json.dump(outputs["action"][0].tolist(), f)
+```
+
+[[autodoc]] Cosmos3OmniModularPipeline
+
+- all
+- __call__
+
 ## CosmosActionCondition
 
 [[autodoc]] CosmosActionCondition
 
 ## Cosmos3OmniPipelineOutput
 
-[[autodoc]] pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput
+[[autodoc]] pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput

src/diffusers/__init__.py

@@ -477,6 +477,8 @@
         [
             "AnimaAutoBlocks",
             "AnimaModularPipeline",
+            "Cosmos3OmniBlocks",
+            "Cosmos3OmniModularPipeline",
             "ErnieImageAutoBlocks",
             "ErnieImageModularPipeline",
             "Flux2AutoBlocks",
@@ -1327,6 +1329,8 @@
         from .modular_pipelines import (
             AnimaAutoBlocks,
             AnimaModularPipeline,
+            Cosmos3OmniBlocks,
+            Cosmos3OmniModularPipeline,
             ErnieImageAutoBlocks,
             ErnieImageModularPipeline,
             Flux2AutoBlocks,

src/diffusers/modular_pipelines/__init__.py

@@ -97,6 +97,10 @@
         "AnimaAutoBlocks",
         "AnimaModularPipeline",
     ]
+    _import_structure["cosmos"] = [
+        "Cosmos3OmniBlocks",
+        "Cosmos3OmniModularPipeline",
+    ]
     _import_structure["ernie_image"] = [
         "ErnieImageAutoBlocks",
         "ErnieImageModularPipeline",
@@ -124,6 +128,7 @@
     else:
         from .anima import AnimaAutoBlocks, AnimaModularPipeline
         from .components_manager import ComponentsManager
+        from .cosmos import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline
         from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline
         from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
         from .flux2 import (

src/diffusers/modular_pipelines/cosmos/__init__.py

@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    DIFFUSERS_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    get_objects_from_module,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+    if not (is_transformers_available() and is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from ...utils import dummy_torch_and_transformers_objects  # noqa F403
+
+    _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+    _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"]
+    _import_structure["modular_pipeline"] = ["Cosmos3OmniModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .modular_blocks_cosmos3 import Cosmos3OmniBlocks
+        from .modular_pipeline import Cosmos3OmniModularPipeline
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )
+
+    for name, value in _dummy_objects.items():
+        setattr(sys.modules[__name__], name, value)