diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 93efbbb5e51e..7c9fcf789913 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -989,11 +989,83 @@ videos = pipe( export_to_video(videos, "cosmos3_modular_transfer_edge.mp4", fps=30, macro_block_size=1) ``` +### Distilled (few-step) text-to-image and image-to-video + +Few-step distilled checkpoints ship a fixed sigma schedule in +`scheduler.config.fixed_step_sampler_config.t_list` and bake classifier-free guidance into +the weights. They are served by the dedicated [`Cosmos3DistilledModularPipeline`] (blocks: +`Cosmos3DistilledBlocks`) — the task-based [`Cosmos3OmniPipeline`] and the base +[`Cosmos3OmniModularPipeline`] do not implement the distilled contract. + +Load a distilled repo and call the pipeline **without** `num_inference_steps` or +`guidance_scale`; `Cosmos3DistilledSetTimestepsStep` reads the fixed step count from the +scheduler config and forces `guidance_scale=1.0`. Because classifier-free guidance is baked +into the weights, `negative_prompt` is not supported (passing one raises an error). + +Prompts follow the same descriptive JSON structure as the non-distilled models, so short text +must be upsampled first — use `--mode text2image` (T2I) or `--mode image2video` (I2V) as +described in [Prompt upsampling](#prompt-upsampling), then pass the JSON via `json.dumps(...)`. + +```python +import json +import torch +from diffusers import Cosmos3DistilledModularPipeline +from diffusers.utils import export_to_video, load_image + +# JSON-upsampled prompt (see "Prompt upsampling" above). +json_prompt = json.load(open("assets/example_t2i_prompt.json")) + +repo = "nvidia/Cosmos3-Super-Text2Image-4Step" +pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# text-to-image (distilled) +videos = pipe( + prompt=json.dumps(json_prompt), + num_frames=1, + height=720, + width=1280, + output="videos", +) +videos[0].save("cosmos3_distilled_t2i.jpg", format="JPEG", quality=85) + +# image-to-video (distilled) — load the I2V repo instead +# JSON-upsampled prompt (see "Prompt upsampling" above); upsampled from the source prompt +# "The right robotic hand picks up the red sphere on the shelf." +json_prompt_i2v = json.load(open("assets/example_i2v_prompt.json")) + +repo_i2v = "nvidia/Cosmos3-Super-Image2Video-4Step" +pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo_i2v, torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.to("cuda") + +image = load_image( + "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" +) +videos = pipe( + prompt=json.dumps(json_prompt_i2v), + image=image, + num_frames=189, + height=720, + width=1280, + output="videos", +) +export_to_video(videos, "cosmos3_distilled_i2v.mp4", fps=24, macro_block_size=1) +``` + [[autodoc]] Cosmos3OmniModularPipeline - all - __call__ +## Cosmos3DistilledModularPipeline + +[[autodoc]] Cosmos3DistilledModularPipeline + +- all +- __call__ + ## CosmosActionCondition [[autodoc]] CosmosActionCondition diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index d7c6ee593a0f..e5c7bc504245 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -489,6 +489,8 @@ [ "AnimaAutoBlocks", "AnimaModularPipeline", + "Cosmos3DistilledBlocks", + "Cosmos3DistilledModularPipeline", "Cosmos3OmniBlocks", "Cosmos3OmniModularPipeline", "ErnieImageAutoBlocks", @@ -1349,6 +1351,8 @@ from .modular_pipelines import ( AnimaAutoBlocks, AnimaModularPipeline, + Cosmos3DistilledBlocks, + Cosmos3DistilledModularPipeline, Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, ErnieImageAutoBlocks, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 25db2ef3bee2..41d8902a053c 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -98,6 +98,8 @@ "AnimaModularPipeline", ] _import_structure["cosmos"] = [ + "Cosmos3DistilledBlocks", + "Cosmos3DistilledModularPipeline", "Cosmos3OmniBlocks", "Cosmos3OmniModularPipeline", ] @@ -128,7 +130,12 @@ else: from .anima import AnimaAutoBlocks, AnimaModularPipeline from .components_manager import ComponentsManager - from .cosmos import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline + from .cosmos import ( + Cosmos3DistilledBlocks, + Cosmos3DistilledModularPipeline, + Cosmos3OmniBlocks, + Cosmos3OmniModularPipeline, + ) from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline from .flux2 import ( diff --git a/src/diffusers/modular_pipelines/cosmos/__init__.py b/src/diffusers/modular_pipelines/cosmos/__init__.py index 20150b299893..38a1b30a421e 100644 --- a/src/diffusers/modular_pipelines/cosmos/__init__.py +++ b/src/diffusers/modular_pipelines/cosmos/__init__.py @@ -22,7 +22,8 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"] - _import_structure["modular_pipeline"] = ["Cosmos3OmniModularPipeline"] + _import_structure["modular_blocks_cosmos3_distilled"] = ["Cosmos3DistilledBlocks"] + _import_structure["modular_pipeline"] = ["Cosmos3DistilledModularPipeline", "Cosmos3OmniModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -32,7 +33,8 @@ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: from .modular_blocks_cosmos3 import Cosmos3OmniBlocks - from .modular_pipeline import Cosmos3OmniModularPipeline + from .modular_blocks_cosmos3_distilled import Cosmos3DistilledBlocks + from .modular_pipeline import Cosmos3DistilledModularPipeline, Cosmos3OmniModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 1e5461429e54..7bf431aa855b 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -5,7 +5,7 @@ from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition -from ...schedulers import UniPCMultistepScheduler +from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam @@ -116,6 +116,11 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=list[int], description="Indexes of conditioned vision latent frames.", ), + OutputParam( + "vision_conditioning_latents", + type_hint=torch.Tensor, + description="Clean encoded vision latents used to re-anchor image conditioning each step.", + ), ] @torch.no_grad() @@ -166,6 +171,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False ).flatten() block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes] + block_state.vision_conditioning_latents = x0_tokens_vision self.set_block_state(state, block_state) return components, state @@ -1237,3 +1243,89 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) ) self.set_block_state(state, block_state) return components, state + + +class Cosmos3DistilledSetTimestepsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Initializes the fixed distilled sampling schedule from the pipeline's `distilled_sigmas` config." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=True), + ConfigSpec(name="distilled_sigmas", default=None), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", required=False, default=None), + InputParam( + name="guidance_scale", + type_hint=float, + default=None, + description=( + "Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the " + "scale is forced to 1.0. Passing a value other than 1.0 raises an error." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."), + OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."), + OutputParam( + "num_inference_steps", + type_hint=int, + description="Resolved number of denoising steps (fixed by the distilled schedule).", + ), + OutputParam( + name="guidance_scale", + type_hint=float, + description="Resolved classifier-free guidance scale (always 1.0 for distilled checkpoints).", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + sigmas = components.config.distilled_sigmas + if not sigmas: + raise ValueError( + "Cosmos3DistilledSetTimestepsStep requires the pipeline config `distilled_sigmas` to be set " + "(populated from the distilled checkpoint's `modular_model_index.json`). Load a distilled Cosmos3 " + "checkpoint or use `Cosmos3OmniModularPipeline` for base checkpoints." + ) + sigmas = [float(s) for s in sigmas] + distilled_steps = len(sigmas) + + if block_state.num_inference_steps is not None and block_state.num_inference_steps != distilled_steps: + raise ValueError( + "This is a distilled checkpoint; the step count is fixed by the pipeline's " + f"`distilled_sigmas` config ({distilled_steps} steps). " + f"`num_inference_steps` must be {distilled_steps} or left unset (got {block_state.num_inference_steps})." + ) + if block_state.guidance_scale is not None and block_state.guidance_scale != 1.0: + raise ValueError( + "This is a distilled checkpoint; classifier-free guidance is baked into the weights. " + f"`guidance_scale` must be 1.0 or left unset (got {block_state.guidance_scale})." + ) + + components.scheduler.set_timesteps(sigmas=sigmas, device=device) + block_state.num_inference_steps = distilled_steps + block_state.guidance_scale = 1.0 + block_state.timesteps = components.scheduler.timesteps + block_state.num_warmup_steps = len(block_state.timesteps) - distilled_steps * components.scheduler.order + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 7d023c7f641c..449acf797d73 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -3,7 +3,7 @@ import torch from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer -from ...schedulers import UniPCMultistepScheduler +from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from ..modular_pipeline import ( BlockState, LoopSequentialPipelineBlocks, @@ -282,6 +282,65 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta return components, block_state +class Cosmos3DistilledVisionLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Updates vision latents after one distilled denoising iteration, re-anchoring conditioned frames." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, description="Noisy vision latents to update."), + InputParam( + name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity." + ), + InputParam( + name="vision_condition_mask", + type_hint=torch.Tensor, + required=True, + description="Mask marking conditioned vision latent frames.", + ), + InputParam( + name="vision_conditioning_latents", + type_hint=torch.Tensor, + default=None, + description="Clean encoded vision latents for re-anchoring conditioned frames.", + ), + InputParam( + name="vision_condition_indexes_for_pack", + type_hint=list, + default=None, + description="Indexes of conditioned vision latent frames; non-empty for image-to-video.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latents = components.scheduler.step( + block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + + # Distilled checkpoints use stochastic (SDE) scheduler steps that re-noise every position. + # Re-anchor conditioned frames to the clean encoded reference after each step. + has_image_condition = bool(block_state.vision_condition_indexes_for_pack) + if has_image_condition and block_state.vision_conditioning_latents is not None: + mask = block_state.vision_condition_mask + reference = block_state.vision_conditioning_latents.to(block_state.latents.dtype) + block_state.latents = mask * reference + (1.0 - mask) * block_state.latents + + return components, block_state + + class Cosmos3SoundLoopSchedulerStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -427,6 +486,27 @@ def description(self) -> str: return "Runs the vision-only Cosmos3 denoising loop." +class Cosmos3DistilledVisionDenoiseStep(Cosmos3DenoiseLoopWrapper): + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3VisionLoopPrepareStep, + Cosmos3LoopDenoiser, + Cosmos3DistilledVisionLoopSchedulerStep, + ] + block_names = ["prepare_vision", "denoiser", "update_vision"] + + @property + def description(self) -> str: + return "Runs the vision-only distilled Cosmos3 denoising loop." + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", Cosmos3OmniTransformer), + ] + + class Cosmos3VisionSoundDenoiseStep(Cosmos3DenoiseLoopWrapper): block_classes = [ Cosmos3VisionLoopPrepareStep, diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index d8c5412078b3..0c53df1e53ec 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -280,6 +280,130 @@ def _tokenize(text: str) -> list[int]: return components, state +class Cosmos3DistilledTextEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Prepares distilled prompt token IDs. Classifier-free guidance is baked into the weights, so " + "`negative_prompt` is not exposed and the unconditional branch is derived from an empty prompt." + ) + + @staticmethod + def _check_inputs(block_state) -> None: + prompt = block_state.prompt + if not isinstance(prompt, str): + raise ValueError( + f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."), + InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."), + InputParam( + name="height", + type_hint=int, + default=None, + description="Height of the generated video or image in pixels.", + ), + InputParam( + name="width", + type_hint=int, + default=None, + description="Width of the generated video or image in pixels.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 system prompt.", + ), + InputParam( + name="add_resolution_template", + type_hint=bool, + default=True, + description="Whether to add resolution metadata to the prompt.", + ), + InputParam( + name="add_duration_template", + type_hint=bool, + default=True, + description="Whether to add duration metadata to the prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("num_frames", type_hint=int, description="Number of frames to generate."), + OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."), + OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."), + OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."), + OutputParam( + "uncond_input_ids", + type_hint=torch.Tensor, + description="Token IDs for the unconditional prompt (empty prompt; guidance is baked in).", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + if block_state.num_frames is None: + block_state.num_frames = 189 + if block_state.height is None: + block_state.height = 720 + if block_state.width is None: + block_state.width = 1280 + + self._check_inputs(block_state) + if components.requires_safety_checker: + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + device = components._execution_device + components.safety_checker.to(device) + try: + if not components.safety_checker.check_text_safety(block_state.prompt): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. " + "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement." + ) + finally: + components.safety_checker.to("cpu") + + # Guidance is baked into distilled weights: the unconditional branch is built from an empty prompt + # (negative_prompt is not a user-facing input) so the downstream text-segment packing contract still holds. + block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( + block_state.prompt, + None, + num_frames=block_state.num_frames, + height=block_state.height, + width=block_state.width, + fps=block_state.fps, + use_system_prompt=block_state.use_system_prompt, + add_resolution_template=block_state.add_resolution_template, + add_duration_template=block_state.add_duration_template, + action_mode=None, + action_view_point=None, + ) + + self.set_block_state(state, block_state) + return components, state + + class Cosmos3ActionTextStep(ModularPipelineBlocks): model_name = "cosmos3-omni" diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py new file mode 100644 index 000000000000..ff2be89f2241 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -0,0 +1,239 @@ +from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + Cosmos3DistilledSetTimestepsStep, + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionPrepareLatentsStep, +) +from .decoders import Cosmos3VideoDecodeStep +from .denoise import Cosmos3DistilledVisionDenoiseStep +from .encoders import ( + Cosmos3DistilledTextEncoderStep, + Cosmos3ImageVaeEncoderStep, + Cosmos3VideoVaeEncoderStep, +) + + +# auto_docstring +class Cosmos3DistilledAutoVaeEncoderStep(ConditionalPipelineBlocks): + """ + Auto VAE conditioning block for distilled Cosmos3. + - Cosmos3VideoVaeEncoderStep runs for the video path. + - Cosmos3ImageVaeEncoderStep runs for the image path. + - when no image or video conditioning is provided, this block is skipped. + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + video (`None`, *optional*): + Reference video for video-to-video conditioning. + condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): + Latent-frame indexes to preserve from the conditioning video. + condition_video_keep (`str`, *optional*, defaults to first): + Which end of a longer conditioning video to use: `first` or `last`. + num_frames (`int`, *optional*): + Number of frames to generate. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + image (`None`, *optional*): + Reference image for image-to-video conditioning. + + Outputs: + x0_tokens_vision (`Tensor`): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`): + Latent-frame indexes fixed by visual conditioning. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep] + block_names = ["video_conditioning", "image_conditioning"] + block_trigger_inputs = ["video", "image"] + default_block_name = None + + def select_block(self, **kwargs) -> str | None: + image = kwargs.get("image") + video = kwargs.get("video") + if image is not None and video is not None: + raise ValueError("Pass either image or video, not both.") + if video is not None: + return "video_conditioning" + if image is not None: + return "image_conditioning" + return None + + @property + def description(self): + return ( + "Auto VAE conditioning block for distilled Cosmos3.\n" + + " - Cosmos3VideoVaeEncoderStep runs for the video path.\n" + + " - Cosmos3ImageVaeEncoderStep runs for the image path.\n" + + " - when no image or video conditioning is provided, this block is skipped." + ) + + +# auto_docstring +class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): + """ + Runs the text-and-vision distilled Cosmos3 denoising workflow. + + Components: + transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Configs: + is_distilled (default: True) distilled_sigmas (default: None) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`): + Number of frames to generate. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*): + The number of denoising steps. + guidance_scale (`float`, *optional*): + Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is + forced to 1.0. Passing a value other than 1.0 raises an error. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionPrepareLatentsStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3DistilledSetTimestepsStep, + Cosmos3DistilledVisionDenoiseStep, + ] + block_names = [ + "prepare_text_segments", + "prepare_vision_latents", + "pack_vision_sequence", + "prepare_vision_denoiser_inputs", + "set_timesteps", + "denoise", + ] + + @property + def description(self): + return "Runs the text-and-vision distilled Cosmos3 denoising workflow." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class Cosmos3DistilledBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes. + + Supported workflows: + - `text2image`: requires `prompt`, `num_frames` + - `text2video`: requires `prompt` + - `image2video`: requires `prompt`, `image` + - `video2video`: requires `prompt`, `video` + + Components: + text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer + (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Configs: + is_distilled (default: True) distilled_sigmas (default: None) + + Inputs: + prompt (`str`): + The text prompt that guides Cosmos3 generation. + num_frames (`int`, *optional*): + Number of frames to generate. + height (`int`, *optional*): + Height of the generated video or image in pixels. + width (`int`, *optional*): + Width of the generated video or image in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 system prompt. + add_resolution_template (`bool`, *optional*, defaults to True): + Whether to add resolution metadata to the prompt. + add_duration_template (`bool`, *optional*, defaults to True): + Whether to add duration metadata to the prompt. + video (`None`, *optional*): + Reference video for video-to-video conditioning. + condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): + Latent-frame indexes to preserve from the conditioning video. + condition_video_keep (`str`, *optional*, defaults to first): + Which end of a longer conditioning video to use: `first` or `last`. + image (`None`, *optional*): + Reference image for image-to-video conditioning. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*): + The number of denoising steps. + guidance_scale (`float`, *optional*): + Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is + forced to 1.0. Passing a value other than 1.0 raises an error. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3DistilledTextEncoderStep, + Cosmos3DistilledAutoVaeEncoderStep, + Cosmos3DistilledVisionCoreDenoiseStep, + Cosmos3VideoDecodeStep, + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True, "num_frames": 1}, + "text2video": {"prompt": True}, + "image2video": {"prompt": True, "image": True}, + "video2video": {"prompt": True, "video": True}, + } + + @property + def description(self): + return "Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index d27fe8310a10..d6c09703c12e 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -116,3 +116,15 @@ def _mask_velocity_predictions(*args, **kwargs): def _apply_video_safety_check(self, *args, **kwargs): return Cosmos3OmniPipeline._apply_video_safety_check(self, *args, **kwargs) + + +class Cosmos3DistilledModularPipeline(Cosmos3OmniModularPipeline): + """ + A ModularPipeline for distilled (few-step) Cosmos 3 omni generation. + + Distilled checkpoints bake classifier-free guidance into the weights and sample on a fixed schedule read from the + pipeline's `distilled_sigmas` config (populated from `modular_model_index.json`), so `guidance_scale` and + `num_inference_steps` are fixed and `negative_prompt` is not supported. + """ + + default_blocks_name = "Cosmos3DistilledBlocks" diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index b03bc6926197..7fdf2c662b1e 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -32,6 +32,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Cosmos3DistilledBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Cosmos3DistilledModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Cosmos3OmniBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py new file mode 100644 index 000000000000..8b44065563b2 --- /dev/null +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -0,0 +1,177 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +import torch + +from diffusers import ModularPipeline +from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline +from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3DistilledSetTimestepsStep +from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep +from diffusers.modular_pipelines.modular_pipeline import PipelineState + +from ...testing_utils import torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +# TODO: move this fixture to `hf-internal-testing/tiny-cosmos3-distilled-modular-pipe` and update the +# repo name here. Hosted on a personal account for now so the PR can be tested. +TINY_DISTILLED_REPO = "yzhautouskay/tiny-cosmos3-distilled-modular-pipe" + + +# text2image / text2video: no visual conditioning, so the auto VAE encoder is skipped. +TEXT_DISTILLED_WORKFLOW = [ + ("text_encoder", "Cosmos3DistilledTextEncoderStep"), + ("denoise.prepare_text_segments", "Cosmos3PrepareTextSegmentsStep"), + ("denoise.prepare_vision_latents", "Cosmos3VisionPrepareLatentsStep"), + ("denoise.pack_vision_sequence", "Cosmos3VisionPackSequenceStep"), + ("denoise.prepare_vision_denoiser_inputs", "Cosmos3VisionDenoiseInputStep"), + ("denoise.set_timesteps", "Cosmos3DistilledSetTimestepsStep"), + ("denoise.denoise", "Cosmos3DistilledVisionDenoiseStep"), + ("decode", "Cosmos3VideoDecodeStep"), +] + +IMAGE_DISTILLED_WORKFLOW = [ + ("text_encoder", "Cosmos3DistilledTextEncoderStep"), + ("vae_encoder", "Cosmos3ImageVaeEncoderStep"), + *TEXT_DISTILLED_WORKFLOW[1:], +] + +VIDEO_DISTILLED_WORKFLOW = [ + ("text_encoder", "Cosmos3DistilledTextEncoderStep"), + ("vae_encoder", "Cosmos3VideoVaeEncoderStep"), + *TEXT_DISTILLED_WORKFLOW[1:], +] + +COSMOS3_DISTILLED_WORKFLOWS = { + "text2image": TEXT_DISTILLED_WORKFLOW, + "text2video": TEXT_DISTILLED_WORKFLOW, + "image2video": IMAGE_DISTILLED_WORKFLOW, + "video2video": VIDEO_DISTILLED_WORKFLOW, +} + + +class TestCosmos3DistilledModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Cosmos3DistilledModularPipeline + pipeline_blocks_class = Cosmos3DistilledBlocks + pretrained_model_name_or_path = TINY_DISTILLED_REPO + + params = frozenset(["prompt", "height", "width", "num_frames"]) + batch_params = frozenset() + optional_params = frozenset(["num_inference_steps", "output_type"]) + output_name = "videos" + expected_workflow_blocks = COSMOS3_DISTILLED_WORKFLOWS + + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipe = super().get_pipeline(components_manager, torch_dtype) + pipe.disable_safety_checker() + return pipe + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "A small robot moves across a table.", + "generator": self.get_generator(seed), + "num_inference_steps": 4, + "height": 32, + "width": 32, + "num_frames": 5, + "output_type": "latent", + } + + def test_save_from_pretrained(self, tmp_path): + base_pipe = self.get_pipeline().to(torch_device) + base_pipe.save_pretrained(str(tmp_path)) + + loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path)) + loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.disable_safety_checker() + loaded_pipe.to(torch_device) + + base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name) + loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert torch.abs(base_output - loaded_output).max() < 1e-3 + + @pytest.mark.skip(reason="Cosmos3 does not support batched prompts.") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason="Cosmos3 does not support batched prompts.") + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason="Cosmos3 does not support multiple videos per prompt.") + def test_num_images_per_prompt(self): + pass + + @pytest.mark.skip(reason="Cosmos3 checkpoints support bfloat16, not float16, inference.") + def test_float16_inference(self): + pass + + +def _fake_distilled_components(sigmas=(1.0, 0.9375, 0.8333333333333334, 0.625)): + config = SimpleNamespace(distilled_sigmas=list(sigmas)) + return SimpleNamespace(_execution_device="cpu", config=config) + + +def test_cosmos3_distilled_vae_encoder_select_block(): + vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"] + assert vae_encoder.select_block(image=None, video=None) is None + assert vae_encoder.select_block(image=object(), video=None) == "image_conditioning" + assert vae_encoder.select_block(image=None, video=object()) == "video_conditioning" + with pytest.raises(ValueError, match="either image or video"): + vae_encoder.select_block(image=object(), video=object()) + + +def test_cosmos3_distilled_set_timesteps_declares_distilled_configs(): + configs = {spec.name: spec.default for spec in Cosmos3DistilledSetTimestepsStep().expected_configs} + assert configs == {"is_distilled": True, "distilled_sigmas": None} + + +def test_cosmos3_distilled_text_encoder_omits_negative_prompt(): + input_names = {inp.name for inp in Cosmos3DistilledTextEncoderStep().inputs} + assert "prompt" in input_names + assert "negative_prompt" not in input_names + + +def test_cosmos3_distilled_text_encoder_requires_str_prompt(): + Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt="a robot")) + with pytest.raises(ValueError, match="`prompt` must be a str"): + Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt=["a robot", "another"])) + + +def test_cosmos3_distilled_set_timesteps_rejects_step_count_override(): + step = Cosmos3DistilledSetTimestepsStep() + state = PipelineState() + state.set("num_inference_steps", 10) + with pytest.raises(ValueError, match="must be 4 or left unset"): + step(_fake_distilled_components(), state) + + +def test_cosmos3_distilled_set_timesteps_rejects_guidance_override(): + step = Cosmos3DistilledSetTimestepsStep() + state = PipelineState() + state.set("guidance_scale", 3.0) + with pytest.raises(ValueError, match="`guidance_scale` must be 1.0"): + step(_fake_distilled_components(), state) + + +def test_cosmos3_distilled_set_timesteps_requires_distilled_sigmas(): + step = Cosmos3DistilledSetTimestepsStep() + components = SimpleNamespace(_execution_device="cpu", config=SimpleNamespace(distilled_sigmas=None)) + with pytest.raises(ValueError, match="distilled_sigmas"): + step(components, PipelineState())