diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 1ff9bc03c382..b02cf7e3739d 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -171,6 +171,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | `Qwen/Qwen-Image-Layered` | Image-to-Image | | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | +| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | ### Feature Matrix diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 2823735d0af4..6353ff8fb7ab 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -40,6 +40,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `Qwen/Qwen-Image-Layered` | Image-to-Image | | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | +| `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically. diff --git a/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml b/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml new file mode 100644 index 000000000000..6ea32adaebf7 --- /dev/null +++ b/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# 1-GPU Cosmos3 text-to-image deployment (base or distilled T2I checkpoints). +# Model: nvidia/Cosmos3-Super-Text2Image or nvidia/Cosmos3-Super-Text2Image-4Step +# Shared by offline examples (--visual_gen_args) and trtllm-serve. +# +# Warmup expresses the deployed workload: warms the 1024x1024 single-frame +# shape instead of the omni default (720p x 189-frame video). Requests should +# pass output_type="image". +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 1 +compilation_config: + resolutions: [[1024, 1024]] + num_frames: [1] diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 69be21fe4880..f5d740841dda 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -13,6 +13,7 @@ Pass the Hub ID or local path via `--model`: - [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) - [`nvidia/Cosmos3-Super`](https://huggingface.co/nvidia/Cosmos3-Super) +- [`nvidia/Cosmos3-Super-Text2Image-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image-4Step) — DMD2-distilled text-to-image: fixed 4-step schedule with classifier-free guidance baked into the weights. Steps/guidance are read from the checkpoint; conflicting request values are rejected. Use with `configs/cosmos3-t2i-1gpu.yaml`. ## Guardrails @@ -36,6 +37,7 @@ See `examples/visual_gen/configs/`: - `cosmos3-nano-1gpu.yaml` — 1 GPU - `cosmos3-super-4gpu.yaml` — 4 GPU, CFG + Ulysses + parallel VAE +- `cosmos3-t2i-1gpu.yaml` — 1 GPU, text-to-image deployments (base or distilled): warms the deployed 1024×1024 single-frame shape instead of the omni video shape. Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`). @@ -70,6 +72,14 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ --output_path output.png +# T2I, distilled 4-step checkpoint (use the T2I config so warmup runs the +# image shape; steps/guidance come from the checkpoint automatically) +python cosmos3.py --model nvidia/Cosmos3-Super-Text2Image-4Step \ + --prompt_file prompts/t2i.json \ + --visual_gen_args ../../configs/cosmos3-t2i-1gpu.yaml \ + --output_type image \ + --output_path output.png + # Inline prompt (--prompt or a JSON file path) python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt "A cute puppy playing with a ball in a park" \ diff --git a/requirements-dev.txt b/requirements-dev.txt index 57d64fb0436c..d5dcc3a2bf7e 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,7 +1,8 @@ -r requirements.txt -# VisualGen LPIPS goldens were recorded with this version. Install it before -# pytest collection so already-imported Diffusers modules match the package files. -diffusers==0.38.0 +# Pin the exact diffusers version VisualGen LPIPS tests run against; install it +# before pytest collection so already-imported Diffusers modules match the +# package files. 0.39.0 matches the runtime floor in requirements.txt. +diffusers==0.39.0 boto3 einops lpips diff --git a/requirements.txt b/requirements.txt index 83ceab0dd1c7..9e10c7b77180 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,8 @@ accelerate>=1.7.0 build colored cuda-python>=13 -diffusers>=0.37.1 +# FlowMatchEuler respects a supplied generator starting in 0.39.0 (huggingface/diffusers#13678). +diffusers>=0.39.0 ftfy lark lazy_loader~=0.5 diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index f5747544946d..114ca97cad22 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -35,10 +35,7 @@ "frame_rate": 24.0, } -# Text-to-image (``output_type="image"``) defaults. Applied by the pipeline when -# the corresponding request field still carries the merged video default, since -# the executor merges a single ``default_generation_params`` dict (the video -# params above) into the request before ``infer()`` runs. +# Text-to-image (``output_type="image"``) defaults; resolved in ``infer()``. COSMOS3_T2I_PARAMS = { "height": 1024, "width": 1024, @@ -48,6 +45,18 @@ "guidance_interval": (400.0, 1000.0), } +# Fields merged by the executor into every request. Mode-dependent values +# remain None until infer() selects the request mode; key membership also +# declares these fields supported during request validation. +COSMOS3_PIPELINE_DEFAULTS = { + **COSMOS3_720P_PARAMS, + "height": None, + "width": None, + "num_inference_steps": None, + "guidance_scale": None, +} + + COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 2dac241f2115..9987a380bedc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -21,7 +21,7 @@ import PIL.Image import torch -from diffusers import AutoencoderKLWan, UniPCMultistepScheduler +from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from transformers import Qwen2Tokenizer @@ -34,8 +34,14 @@ from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS +from .defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_EXTRA_SPECS, + COSMOS3_PIPELINE_DEFAULTS, + COSMOS3_T2I_PARAMS, +) from .guardrails import check_video_safety, download_guardrail_checkpoint +from .sampling import Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer @@ -60,6 +66,7 @@ "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Image2Video", "nvidia/Cosmos3-Super-Text2Image", + "nvidia/Cosmos3-Super-Text2Image-4Step", ], doc="Cosmos3 Omnimodal world models.", ) @@ -68,6 +75,9 @@ def __init__(self, pipeline_config): primary_pretrained_config = pipeline_config.primary_pretrained_config self.audio_gen = False self.action_gen = False + # Pre-load placeholder; load_standard_components derives the real + # policy from the checkpoint's scheduler via from_scheduler(). + self.sampling = Cosmos3SamplingPolicy() if getattr( primary_pretrained_config, "audio_gen", @@ -139,22 +149,15 @@ def load_standard_components( if PipelineComponent.SCHEDULER not in skip_components: logger.info("Loading scheduler...") - self.scheduler = UniPCMultistepScheduler.from_pretrained( - checkpoint_dir, - subfolder=PipelineComponent.SCHEDULER, - ) - # Snapshot the checkpoint scheduler config so the scheduler can be - # rebuilt at request time when a mode-specific ``flow_shift`` is - # needed (T2I uses shift=3.0; T2V/I2V keep the checkpoint default). - self._base_scheduler_config = self.scheduler.config - self._engine_init_flow_shift = float( - getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0 - ) - self._current_flow_shift = self._engine_init_flow_shift + # The scheduler class comes from the checkpoint: UniPC for base + # checkpoints, FlowMatchEuler (fixed stochastic schedule) for + # distilled ones. The policy holds the derived immutable facts. + self.scheduler = load_scheduler(checkpoint_dir) + self.sampling = Cosmos3SamplingPolicy.from_scheduler(self.scheduler) if self.audio_gen: - # Separate instance so video and audio scheduler states don't collide - # (UniPC mutates internal correction buffers on every .step() call). - self.audio_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + # Separate instance so video and audio scheduler states don't + # collide (schedulers mutate internal state on every .step()). + self.audio_scheduler = type(self.scheduler).from_config(self.scheduler.config) # Re-check the env var in case it was changed after initialization like in unit tests. guardrails_disabled = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -185,23 +188,6 @@ def load_standard_components( self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - def _set_flow_shift(self, target_shift: float) -> None: - """Rebuild the UniPC scheduler with ``flow_shift=target_shift`` if needed. - - T2I uses ``flow_shift=3.0`` while T2V/I2V use the checkpoint default. - ``self._current_flow_shift`` is tracked explicitly so a prior T2I rebuild - does not leak into a subsequent video request. - """ - if not hasattr(self, "_base_scheduler_config"): - return - target = float(target_shift) - if target == float(self._current_flow_shift): - return - self.scheduler = UniPCMultistepScheduler.from_config( - self._base_scheduler_config, flow_shift=target - ) - self._current_flow_shift = target - @property def default_warmup_resolutions(self): return [(720, 1280)] @@ -210,15 +196,26 @@ def default_warmup_resolutions(self): def default_warmup_num_frames(self): return [189] + @property + def default_warmup_steps(self): + # Distilled checkpoints only run their fixed schedule length. + return self.sampling.num_steps(super().default_warmup_steps) + @property def default_generation_params(self): - return dict(COSMOS3_720P_PARAMS) + return {**COSMOS3_PIPELINE_DEFAULTS, **self.sampling.generation_default_overrides()} @property def extra_param_specs(self): return dict(COSMOS3_EXTRA_SPECS) def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: + # Checkpoint-aware guidance: distilled defaults carry a concrete 1.0; + # base defaults leave it None ("by mode") — warmup runs the video mode. + defaults = self.default_generation_params + guidance_scale = defaults["guidance_scale"] + if guidance_scale is None: + guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] with torch.no_grad(): self.forward( prompt="warmup", @@ -227,51 +224,29 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N width=width, num_frames=num_frames, num_inference_steps=steps, - guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + guidance_scale=guidance_scale, seed=42, - max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + max_sequence_length=defaults["max_sequence_length"], use_guardrails=False, image=None, enable_audio=False, ) - @staticmethod - def _resolve_t2i_default(merged_value, video_default, t2i_default): - """Pick the T2I default when the field still carries the merged video default. - - The executor merges a single ``default_generation_params`` dict (the - video params) into the request before ``infer()``, so an unspecified - field arrives equal to its video default. For T2I we substitute the - T2I default in that case while honoring any explicit user override. - """ - return t2i_default if merged_value == video_default else merged_value - def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") is_t2i = str(output_type).lower() == "image" - height = req.params.height - width = req.params.width - num_inference_steps = req.params.num_inference_steps - guidance_scale = req.params.guidance_scale - if is_t2i: - height = self._resolve_t2i_default( - height, COSMOS3_720P_PARAMS["height"], COSMOS3_T2I_PARAMS["height"] - ) - width = self._resolve_t2i_default( - width, COSMOS3_720P_PARAMS["width"], COSMOS3_T2I_PARAMS["width"] - ) - num_inference_steps = self._resolve_t2i_default( - num_inference_steps, - COSMOS3_720P_PARAMS["num_inference_steps"], - COSMOS3_T2I_PARAMS["num_inference_steps"], - ) - guidance_scale = self._resolve_t2i_default( - guidance_scale, - COSMOS3_720P_PARAMS["guidance_scale"], - COSMOS3_T2I_PARAMS["guidance_scale"], - ) + # None = unset; resolve by mode exactly once. Non-None values pass through. + mode_params = COSMOS3_T2I_PARAMS if is_t2i else COSMOS3_720P_PARAMS + + def resolved(value, field_name): + return value if value is not None else mode_params[field_name] + + height = resolved(req.params.height, "height") + width = resolved(req.params.width, "width") + num_inference_steps = resolved(req.params.num_inference_steps, "num_inference_steps") + guidance_scale = resolved(req.params.guidance_scale, "guidance_scale") return self.forward( prompt=req.prompt, @@ -634,6 +609,15 @@ def forward( enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, ): + """Run one generation. ``infer()`` is the resolved entry point. + + Production requests arrive through ``infer()`` with fully resolved + values; the signature defaults are the base-checkpoint *video* table + values for direct internal callers. ``forward()`` cannot tell a + signature default from an explicit argument, so on distilled + checkpoints (which fix steps/guidance and reject anything else) direct + callers must pass checkpoint-valid sampling values. + """ pipeline_start = time.time() timer = CudaPhaseTimer() timer.mark_pre_start() @@ -643,7 +627,20 @@ def forward( # Text-to-image mode: same checkpoint/forward path as T2V, but a single # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG # guidance interval, and an image (rather than video) output. - is_t2i = str(output_type).lower() == "image" + output_type = str(output_type).lower() + if output_type not in ("video", "image"): + raise ValueError(f"output_type must be 'video' or 'image', got {output_type!r}.") + is_t2i = output_type == "image" + + self.sampling.validate_request(num_inference_steps, guidance_scale) + + if image is not None and self.sampling.is_distilled: + raise ValueError( + "Image-conditioned generation is not supported on distilled Cosmos3 " + "checkpoints yet: the stochastic scheduler re-noises the conditioned " + "frame at every step, and this pipeline does not re-anchor it per step." + ) + guidance_interval = None if is_t2i: if image is not None: @@ -653,11 +650,15 @@ def forward( num_frames = 1 enable_audio = False guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] - self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) + self.scheduler = self.sampling.set_flow_shift( + self.scheduler, COSMOS3_T2I_PARAMS["flow_shift"] + ) else: # Restore the checkpoint flow_shift in case a prior T2I request # rebuilt the scheduler with shift=3.0. - self._set_flow_shift(getattr(self, "_engine_init_flow_shift", 1.0)) + self.scheduler = self.sampling.set_flow_shift( + self.scheduler, self.sampling.checkpoint_flow_shift + ) if isinstance(prompt, str): prompt = [prompt] @@ -781,7 +782,7 @@ def forward( video_shape = (T_latent, H_latent, W_latent) # 3. Set up scheduler - self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self.sampling.set_timesteps(self.scheduler, num_inference_steps, device=self.device) # 3b. Audio noise init — latent length matches diffusers Cosmos3OmniPipeline.prepare_latents. do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") @@ -798,7 +799,9 @@ def forward( dtype=latents.dtype, ) # Audio uses the same scheduler type/config as video. - self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) + self.sampling.set_timesteps( + self.audio_scheduler, num_inference_steps, device=self.device + ) # 4. Build forward_fn for the denoise loop def forward_fn( @@ -861,6 +864,7 @@ def forward_fn( extra_cfg_tensors=extra_cfg_tensors, extra_streams=extra_streams, guidance_interval=guidance_interval, + scheduler_step_kwargs=self.sampling.scheduler_step_kwargs(generator), ) if extra_streams is not None: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py new file mode 100644 index 000000000000..932ac974658a --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Per-checkpoint sampling policy for Cosmos3. + +Exactly two recipes are supported, read from the checkpoint's +``scheduler/scheduler_config.json``: + +* ``UniPCMultistepScheduler`` without fixed sigmas — base checkpoints: + request tables drive steps/guidance; T2I rebuilds with ``flow_shift=3.0``. +* ``FlowMatchEulerDiscreteScheduler`` with ``stochastic_sampling`` enabled + and a nonempty ``fixed_step_sampler_config.t_list`` — distilled + checkpoints: the step count is locked to the schedule, classifier-free + guidance is baked into the weights (scale 1.0), and every step draws + seeded SDE noise. + +The pipeline owns the scheduler instances; :class:`Cosmos3SamplingPolicy` is +an immutable value object of config-derived facts whose methods take +schedulers as arguments. +""" + +import json +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Optional + +from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler + +from tensorrt_llm.logger import logger + +# Distilled checkpoints bake classifier-free guidance into the weights; the +# only valid scale is 1.0 ("off": a single conditional forward per step). +DISTILLED_GUIDANCE_SCALE = 1.0 + + +def _config_get(config: Any, key: str, default: Any = None) -> Any: + """Fetch a key from a plain dict, a diffusers FrozenDict, or a config object.""" + if isinstance(config, Mapping): + return config.get(key, default) + return getattr(config, key, default) + + +def _resolve_distilled_sigmas(scheduler_config: Any) -> "tuple[float, ...] | None": + """``fixed_step_sampler_config.t_list`` as floats, or None (base checkpoints).""" + fixed_step_cfg = _config_get(scheduler_config, "fixed_step_sampler_config") + t_list = _config_get(fixed_step_cfg, "t_list") if fixed_step_cfg else None + if not t_list: + return None + return tuple(float(sigma) for sigma in t_list) + + +def load_scheduler(checkpoint_dir: str, subfolder: str = "scheduler") -> Any: + """Instantiate the scheduler class the checkpoint declares. + + Base checkpoints declare ``UniPCMultistepScheduler`` (a missing declaration + also resolves to UniPC, preserving pre-declaration Cosmos3 behavior); + distilled ones declare ``FlowMatchEulerDiscreteScheduler``. An explicitly + unknown declaration is a load-time error — silently substituting UniPC + would sample the checkpoint with the wrong integrator. + """ + config_path = os.path.join(checkpoint_dir, subfolder, "scheduler_config.json") + class_name = "" + if os.path.exists(config_path): + with open(config_path) as f: + class_name = json.load(f).get("_class_name", "") + if class_name == "FlowMatchEulerDiscreteScheduler": + scheduler_cls = FlowMatchEulerDiscreteScheduler + elif class_name in ("", None, "UniPCMultistepScheduler"): + scheduler_cls = UniPCMultistepScheduler + else: + raise ValueError( + f"Unsupported Cosmos3 scheduler class {class_name!r}; supported: " + "UniPCMultistepScheduler (base), FlowMatchEulerDiscreteScheduler (distilled)." + ) + return scheduler_cls.from_pretrained(checkpoint_dir, subfolder=subfolder) + + +@dataclass(frozen=True) +class Cosmos3SamplingPolicy: + """Immutable sampling facts of a loaded Cosmos3 checkpoint. + + Construct via :meth:`from_scheduler`, which validates the recipe at load + time. A default-constructed policy (all fields ``None``) is the explicit + pre-load placeholder the pipeline holds before its scheduler exists: not + distilled, no flow-shift rebuild capability (``set_flow_shift`` is a + no-op), and replaced by ``from_scheduler`` when components load. + + Methods take scheduler instances as arguments; the current flow shift is + read from the supplied scheduler's config rather than tracked here. + """ + + # Fixed distilled schedule (t_list); None for base checkpoints. + fixed_sigmas: "tuple[float, ...] | None" = None + # Checkpoint scheduler config, kept for flow-shift rebuilds (UniPC only). + unipc_base_config: Optional[Any] = None + + @classmethod + def from_scheduler(cls, scheduler: Any) -> "Cosmos3SamplingPolicy": + """Derive the policy from a loaded scheduler's config. + + Valid recipes: UniPC without fixed sigmas (base) and stochastic + FlowMatchEuler with a nonempty ``fixed_step_sampler_config.t_list`` + (distilled); anything else fails here, at load time. + """ + fixed_sigmas = _resolve_distilled_sigmas(scheduler.config) + is_unipc = isinstance(scheduler, UniPCMultistepScheduler) + is_flow_match = isinstance(scheduler, FlowMatchEulerDiscreteScheduler) + + if ( + _config_get(scheduler.config, "fixed_step_requires_explicit_sigmas", False) + and fixed_sigmas is None + ): + raise ValueError( + "Malformed distilled checkpoint: the scheduler config declares " + "fixed_step_requires_explicit_sigmas but carries no usable " + "fixed_step_sampler_config.t_list." + ) + + if is_unipc and fixed_sigmas is None: + return cls(fixed_sigmas=None, unipc_base_config=scheduler.config) + + if is_flow_match and fixed_sigmas is not None: + if not _config_get(scheduler.config, "stochastic_sampling", False): + raise ValueError( + "Unsupported Cosmos3 sampling recipe: FlowMatchEulerDiscreteScheduler " + "declares a fixed step schedule without stochastic_sampling. The " + "distilled recipe draws SDE noise at every step; running the schedule " + "as an ODE would sample the checkpoint incorrectly." + ) + fixed_step_cfg = _config_get(scheduler.config, "fixed_step_sampler_config") + sample_type = _config_get(fixed_step_cfg, "sample_type") + if sample_type is not None and sample_type != "sde": + raise ValueError( + "Unsupported Cosmos3 sampling recipe: fixed_step_sampler_config " + f"declares sample_type={sample_type!r}; only 'sde' is supported." + ) + logger.info( + f"Distilled Cosmos3 checkpoint: fixed {len(fixed_sigmas)}-step schedule " + f"{list(fixed_sigmas)}, classifier-free guidance baked in." + ) + return cls(fixed_sigmas=fixed_sigmas, unipc_base_config=None) + + raise ValueError( + f"Unsupported Cosmos3 sampling recipe: {type(scheduler).__name__} with " + f"fixed sigmas {'present' if fixed_sigmas is not None else 'absent'}. " + "Supported: UniPCMultistepScheduler without fixed sigmas (base), " + "stochastic FlowMatchEulerDiscreteScheduler with " + "fixed_step_sampler_config.t_list (distilled)." + ) + + @property + def is_distilled(self) -> bool: + return self.fixed_sigmas is not None + + def generation_default_overrides(self) -> dict: + """Checkpoint-mandated overrides of the table generation defaults. + + Merged over ``COSMOS3_720P_PARAMS`` by the pipeline's + ``default_generation_params``, so executor-merged requests arrive + carrying the checkpoint's true defaults. + """ + if not self.is_distilled: + return {} + return { + "num_inference_steps": len(self.fixed_sigmas), + "guidance_scale": DISTILLED_GUIDANCE_SCALE, + } + + def num_steps(self, default: int) -> int: + """The only step count this policy can run: fixed for distilled, else ``default``.""" + return len(self.fixed_sigmas) if self.is_distilled else default + + def validate_request( + self, num_inference_steps: Optional[int], guidance_scale: Optional[float] + ) -> None: + """Reject sampling parameters incompatible with a distilled checkpoint.""" + if not self.is_distilled: + return + distilled_steps = len(self.fixed_sigmas) + if num_inference_steps is not None and num_inference_steps != distilled_steps: + raise ValueError( + "This is a distilled Cosmos3 checkpoint; the step count is fixed by the " + f"scheduler's fixed_step_sampler_config.t_list ({distilled_steps} steps). " + f"num_inference_steps must be {distilled_steps} or left unset " + f"(got {num_inference_steps})." + ) + if guidance_scale is not None and float(guidance_scale) != DISTILLED_GUIDANCE_SCALE: + raise ValueError( + "This is a distilled Cosmos3 checkpoint; classifier-free guidance is baked " + f"into the weights. guidance_scale must be {DISTILLED_GUIDANCE_SCALE} or " + f"left unset (got {guidance_scale})." + ) + + def set_timesteps(self, scheduler: Any, num_inference_steps: int, device: Any) -> None: + """Program a scheduler for one generation: fixed sigmas or a step count.""" + if self.is_distilled: + scheduler.set_timesteps(sigmas=list(self.fixed_sigmas), device=device) + else: + scheduler.set_timesteps(num_inference_steps, device=device) + + def scheduler_step_kwargs(self, generator: Any) -> dict: + """Extra kwargs each ``scheduler.step()`` call requires. + + The distilled FlowMatchEuler scheduler is stochastic: every step draws + SDE noise, which must come from the request-seeded ``generator`` — + otherwise it comes from the process-global RNG, breaking seed + reproducibility and diverging the replicated latents across ranks + (each rank's global RNG state is independent). UniPC steps are + deterministic and accept no ``generator`` argument, so base + checkpoints pass nothing. + """ + if self.is_distilled: + return {"generator": generator} + return {} + + @property + def checkpoint_flow_shift(self) -> float: + """The flow shift the checkpoint shipped with (UniPC only; 1.0 otherwise).""" + if self.unipc_base_config is None: + return 1.0 + return float(_config_get(self.unipc_base_config, "flow_shift", 1.0) or 1.0) + + def set_flow_shift(self, scheduler: Any, target_shift: Optional[float]) -> Any: + """Return ``scheduler`` rebuilt with ``flow_shift=target_shift`` if needed. + + The current shift is read from the supplied scheduler's own config, so + no tracking state exists to diverge. Structural no-op for distilled + checkpoints (no UniPC base config) and for ``target_shift=None``. + """ + if target_shift is None or self.unipc_base_config is None: + return scheduler + target_shift = float(target_shift) + current_shift = float(_config_get(scheduler.config, "flow_shift", 1.0) or 1.0) + if current_shift == target_shift: + return scheduler + return UniPCMultistepScheduler.from_config(self.unipc_base_config, flow_shift=target_shift) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index b3df78a06ec1..1a85cf2b5784 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -15,7 +15,7 @@ import math from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional, Tuple, TypeVar import torch import torch.nn as nn @@ -34,6 +34,27 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig +# Some Cosmos3OmniTransformer checkpoint configs omit these fields; the values +# match what other conversions carry explicitly. +PRETRAINED_CONFIG_COMPAT_DEFAULTS = { + "position_embedding_type": "unified_3d_mrope", + "max_position_embeddings": 262144, + "temporal_compression_factor_sound": 1, +} + + +_PretrainedConfigT = TypeVar("_PretrainedConfigT") + + +def apply_pretrained_config_compat_defaults( + pretrained_config: _PretrainedConfigT, +) -> _PretrainedConfigT: + """Fill missing schema fields in place (idempotent); returns the config.""" + for key, value in PRETRAINED_CONFIG_COMPAT_DEFAULTS.items(): + if getattr(pretrained_config, key, None) is None: + setattr(pretrained_config, key, value) + return pretrained_config + class Qwen3VLTextRMSNorm(nn.Module): def __init__( @@ -702,7 +723,7 @@ def forward( class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) - pretrained_config = model_config.pretrained_config + pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) self.action_gen = getattr(pretrained_config, "action_gen", False) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 13cbe77bdf12..08c8acae8b35 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -978,16 +978,22 @@ def _scheduler_step( timestep, scheduler, extra_stream_schedulers, + scheduler_step_kwargs=None, ): """Execute scheduler step for all streams.""" + step_kwargs = scheduler_step_kwargs or {} t_start = time.time() - latents = scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] + latents = scheduler.step(noise_pred, timestep, latents, return_dict=False, **step_kwargs)[0] # Step schedulers for extra streams for name, noise_extra in extra_noise_preds.items(): if name in extra_stream_schedulers: extra_stream_latents[name] = extra_stream_schedulers[name].step( - noise_extra, timestep, extra_stream_latents[name], return_dict=False + noise_extra, + timestep, + extra_stream_latents[name], + return_dict=False, + **step_kwargs, )[0] t_sched = time.time() - t_start @@ -1010,6 +1016,7 @@ def denoise( boundary_timestep: Optional[float] = None, guidance_interval: Optional[Tuple[float, float]] = None, post_step_fn: Optional[Callable] = None, + scheduler_step_kwargs: Optional[Dict[str, Any]] = None, ): """Execute denoising loop with optional CFG parallel and TeaCache support. @@ -1045,6 +1052,8 @@ def denoise( post_step_fn: Optional callable applied to latents after each scheduler step. Signature: post_step_fn(latents) -> latents Use for constraints that must hold throughout denoising. + scheduler_step_kwargs: Extra keyword arguments forwarded to every + scheduler's ``step()`` call. Returns: Single latents if no extra_streams @@ -1171,6 +1180,7 @@ def denoise( t, scheduler, extra_stream_schedulers, + scheduler_step_kwargs=scheduler_step_kwargs, ) if post_step_fn is not None: diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ce83f862326f..9648d26f87ce 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -283,11 +283,16 @@ def extra_param_specs(self) -> Dict[str, "ExtraParamSchema"]: @property def default_params(self) -> "VisualGenParams": - """Returns a ``VisualGenParams`` with all defaults resolved for the loaded pipeline. + """Returns a ``VisualGenParams`` with the loaded pipeline's defaults. Universal fields (height, width, etc.) are filled from the - pipeline's defaults. All declared ``extra_params`` keys are - included with their defaults (``None`` for params without one). + pipeline's defaults. Pipelines with mode-dependent defaults + (e.g. Cosmos3, where text-to-image and video requests use + different resolutions) leave such fields as ``None``; they are + resolved per request from the output mode, so ``None`` here + means "the mode's default", not "unset". All declared + ``extra_params`` keys are included with their defaults + (``None`` for params without one). Use this to inspect what the model will use, then modify and pass to ``generate()``:: diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index ffcb9c7863d4..8f7b2c0337bc 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -2062,3 +2062,52 @@ def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, ) assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + + +def test_cosmos3_t2i_4step_example(_visual_gen_deps, llm_root, llm_venv): + """Run the distilled T2I checkpoint through the recommended invocation. + + Validates the documented deployment for ``Cosmos3-Super-Text2Image-4Step``: + the example script with ``configs/cosmos3-t2i-1gpu.yaml`` (T2I warmup + shapes) and ``--output_type image``. Steps/guidance come from the + checkpoint's fixed distilled schedule; the run must produce an image. + """ + model_path = _lpips_model_path("Cosmos3-Super-Text2Image-4Step") + _skip_if_missing(model_path, "Cosmos3-Super-Text2Image-4Step checkpoint", is_dir=True) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_t2i_4step_example" + ) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "cosmos3_t2i_4step_output.png") + if os.path.exists(output_path): + os.remove(output_path) + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "cosmos3-t2i-1gpu.yaml" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + assert os.path.isfile(config_path), f"Config not found: {config_path}" + + _venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--prompt", + "A ceramic teapot pouring steaming tea into a cup, morning window light", + "--output_type", + "image", + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty image at {output_path}" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1b3a17bdf504..4fa2a68f9aec 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -264,6 +264,7 @@ l0_b200: tests: - unittest/llmapi/test_llm_quant.py # 3.5 mins on B200 - unittest/disaggregated/test_openai_server_info.py + - examples/visual_gen/test_visual_gen.py::test_cosmos3_t2i_4step_example TIMEOUT (30) - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/visual_gen/conftest.py b/tests/unittest/_torch/visual_gen/conftest.py new file mode 100644 index 000000000000..f6d6187b2e21 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/conftest.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared pytest configuration for the VisualGen unit tests.""" + +from collections.abc import Iterator + +import pytest + + +@pytest.fixture(scope="module") +def disable_cosmos3_guardrails() -> Iterator[None]: + """Disable Cosmos3 guardrails for the requesting module, leak-free. + + Patches both the environment variable (re-read by + ``load_standard_components`` on every call) and the pipeline module's + derived global (assigned by that same function), so teardown restores + both. Opt in per module via ``pytest.mark.usefixtures``. + """ + import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as pipe_mod + + patcher = pytest.MonkeyPatch() + patcher.setenv("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + patcher.setattr(pipe_mod, "TRTLLM_DISABLE_COSMOS3_GUARDRAILS", True) + yield + patcher.undo() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py new file mode 100644 index 000000000000..955eec6c29ab --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Cosmos3 sampling policy (base vs distilled checkpoints) +and its pipeline wiring: scheduler loading, recipe validation, generation +defaults, mode resolution, and the guidance-1.0 denoise-loop contract.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler + +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_T2I_PARAMS, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import ( + DISTILLED_GUIDANCE_SCALE, + Cosmos3SamplingPolicy, + load_scheduler, +) +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline + +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] + +# The relevant subset of the 4-Step checkpoint's scheduler config +# (values verbatim; keys unrelated to distilled detection omitted). +DISTILLED_SIGMAS = (1.0, 0.9375, 0.8333333333333334, 0.625) +DISTILLED_SCHEDULER_CONFIG = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "num_train_timesteps": 1000, + "shift": 1.0, + "stochastic_sampling": True, + "use_karras_sigmas": False, + "fixed_step_requires_explicit_sigmas": True, + "fixed_step_sampler_config": { + "sample_type": "sde", + "t_list": list(DISTILLED_SIGMAS), + }, +} + +UNIPC_SCHEDULER_CONFIG = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "solver_order": 2, +} + +SKIP_NON_SCHEDULER = ["text_tokenizer", "tokenizer", "vae", "sound_tokenizer"] + + +def _write_scheduler_config(checkpoint_dir: Path, config: dict) -> None: + scheduler_dir = checkpoint_dir / "scheduler" + scheduler_dir.mkdir(parents=True, exist_ok=True) + with open(scheduler_dir / "scheduler_config.json", "w") as f: + json.dump(config, f) + + +def _distilled_policy() -> Cosmos3SamplingPolicy: + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + return Cosmos3SamplingPolicy.from_scheduler(scheduler) + + +def _base_scheduler() -> UniPCMultistepScheduler: + return UniPCMultistepScheduler.from_config(UNIPC_SCHEDULER_CONFIG) + + +def _base_policy() -> Cosmos3SamplingPolicy: + return Cosmos3SamplingPolicy.from_scheduler(_base_scheduler()) + + +def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: + """A pipeline instance without heavyweight __init__; ``rank``/``dtype``/ + ``device`` are BasePipeline properties and must not be set here.""" + pipeline = object.__new__(Cosmos3OmniMoTPipeline) + defaults = dict( + audio_gen=False, + action_gen=False, + sampling=Cosmos3SamplingPolicy(), + ) + defaults.update(attrs) + for key, value in defaults.items(): + setattr(pipeline, key, value) + return pipeline + + +def _fake_request(output_type: str = "video", **param_overrides) -> SimpleNamespace: + """A DiffusionRequest look-alike with executor-merged (None = unset) params.""" + params = SimpleNamespace( + height=None, + width=None, + num_inference_steps=None, + guidance_scale=None, + num_frames=COSMOS3_720P_PARAMS["num_frames"], + max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + frame_rate=COSMOS3_720P_PARAMS["frame_rate"], + seed=0, + negative_prompt=None, + image=None, + extra_params={"output_type": output_type}, + ) + for key, value in param_overrides.items(): + setattr(params, key, value) + return SimpleNamespace(prompt="x", params=params) + + +class TestSchedulerLoading: + def test_flow_match_declared(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + assert isinstance(load_scheduler(str(tmp_path)), FlowMatchEulerDiscreteScheduler) + + def test_unipc_declared(self, tmp_path): + _write_scheduler_config(tmp_path, UNIPC_SCHEDULER_CONFIG) + assert isinstance(load_scheduler(str(tmp_path)), UniPCMultistepScheduler) + + def test_missing_class_name_defaults_to_unipc(self, tmp_path): + config = {k: v for k, v in UNIPC_SCHEDULER_CONFIG.items() if k != "_class_name"} + _write_scheduler_config(tmp_path, config) + assert isinstance(load_scheduler(str(tmp_path)), UniPCMultistepScheduler) + + def test_unknown_class_name_raises(self, tmp_path): + """Silently substituting UniPC for an unknown declared scheduler would + sample the checkpoint with the wrong integrator.""" + _write_scheduler_config(tmp_path, {**UNIPC_SCHEDULER_CONFIG, "_class_name": "DDIM"}) + with pytest.raises(ValueError, match="DDIM"): + load_scheduler(str(tmp_path)) + + +class TestPolicyFacts: + def test_is_distilled(self): + assert _distilled_policy().is_distilled + assert not _base_policy().is_distilled + assert not Cosmos3SamplingPolicy().is_distilled + + def test_diffusers_retains_unexpected_config_keys(self): + """Canary: diffusers must keep the unexpected fixed_step_sampler_config + key in scheduler.config; if an upgrade drops it, distilled detection + silently breaks.""" + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler) + assert policy.fixed_sigmas == DISTILLED_SIGMAS + assert scheduler.config.stochastic_sampling is True + + def test_sigma_values_coerced_to_floats(self): + config = { + **DISTILLED_SCHEDULER_CONFIG, + "fixed_step_sampler_config": {"t_list": [1, "0.5"]}, + } + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + assert Cosmos3SamplingPolicy.from_scheduler(scheduler).fixed_sigmas == (1.0, 0.5) + + def test_generation_default_overrides(self): + assert _distilled_policy().generation_default_overrides() == { + "num_inference_steps": 4, + "guidance_scale": DISTILLED_GUIDANCE_SCALE, + } + assert _base_policy().generation_default_overrides() == {} + + def test_num_steps(self): + assert _distilled_policy().num_steps(2) == 4 + assert _base_policy().num_steps(2) == 2 + + def test_checkpoint_flow_shift(self): + assert _base_policy().checkpoint_flow_shift == 1.0 + assert _distilled_policy().checkpoint_flow_shift == 1.0 # no UniPC config + + def test_scheduler_step_kwargs(self): + generator = torch.Generator().manual_seed(7) + assert _distilled_policy().scheduler_step_kwargs(generator) == {"generator": generator} + assert _base_policy().scheduler_step_kwargs(generator) == {} + + +class TestMalformedRecipeValidation: + """Only two recipes are valid; everything else must fail at load.""" + + @pytest.mark.parametrize("broken_fixed_step", [None, {}, {"t_list": []}]) + def test_required_sigmas_missing_raises(self, broken_fixed_step): + config = {k: v for k, v in DISTILLED_SCHEDULER_CONFIG.items()} + config.pop("fixed_step_sampler_config") + if broken_fixed_step is not None: + config["fixed_step_sampler_config"] = broken_fixed_step + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="fixed_step_requires_explicit_sigmas"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_t_list_on_unipc_raises(self): + """UniPC cannot honor the distilled policy (no seeded step noise, + no baked-in guidance) even though its set_timesteps accepts sigmas.""" + config = { + **UNIPC_SCHEDULER_CONFIG, + "fixed_step_sampler_config": DISTILLED_SCHEDULER_CONFIG["fixed_step_sampler_config"], + } + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_unipc_with_declared_requirement_but_no_sigmas_raises(self): + config = {**UNIPC_SCHEDULER_CONFIG, "fixed_step_requires_explicit_sigmas": True} + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="fixed_step_requires_explicit_sigmas"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_unipc_with_flag_and_sigmas_gets_unsupported_error(self): + config = { + **UNIPC_SCHEDULER_CONFIG, + "fixed_step_requires_explicit_sigmas": True, + "fixed_step_sampler_config": DISTILLED_SCHEDULER_CONFIG["fixed_step_sampler_config"], + } + scheduler = UniPCMultistepScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_flow_match_without_sigmas_raises(self): + config = {k: v for k, v in DISTILLED_SCHEDULER_CONFIG.items()} + config.pop("fixed_step_sampler_config") + config.pop("fixed_step_requires_explicit_sigmas") + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="Unsupported Cosmos3 sampling recipe"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + @pytest.mark.parametrize("stochastic", [False, "absent"]) + def test_non_stochastic_fixed_schedule_raises(self, stochastic): + """The distilled policy assumes SDE noise every step (seeded generator); + an ODE fixed-step recipe must not silently load as distilled.""" + config = dict(DISTILLED_SCHEDULER_CONFIG) + if stochastic == "absent": + config.pop("stochastic_sampling") # diffusers defaults it to False + else: + config["stochastic_sampling"] = stochastic + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="stochastic_sampling"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + def test_declared_non_sde_sample_type_raises(self): + config = { + **DISTILLED_SCHEDULER_CONFIG, + "fixed_step_sampler_config": { + "sample_type": "ode", + "t_list": list(DISTILLED_SIGMAS), + }, + } + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + + with pytest.raises(ValueError, match="sample_type"): + Cosmos3SamplingPolicy.from_scheduler(scheduler) + + +class TestValidateRequest: + @pytest.mark.parametrize("steps", [None, 4]) + @pytest.mark.parametrize("guidance", [None, 1, 1.0]) + def test_valid_values_pass(self, steps, guidance): + _distilled_policy().validate_request(steps, guidance) + + @pytest.mark.parametrize("bad_steps", [1, 10, 35, 50, 100]) + def test_explicit_steps_mismatch_raises(self, bad_steps): + with pytest.raises(ValueError, match="distilled"): + _distilled_policy().validate_request(bad_steps, None) + + @pytest.mark.parametrize("bad_guidance", [0.5, 3.5, 6.0, 7.0]) + def test_explicit_guidance_mismatch_raises(self, bad_guidance): + with pytest.raises(ValueError, match="distilled"): + _distilled_policy().validate_request(None, bad_guidance) + + def test_base_policy_accepts_anything(self): + _base_policy().validate_request(17, 5.5) + _base_policy().validate_request(None, None) + + +class TestFlowShift: + def test_unipc_rebuilds_on_change(self): + policy = _base_policy() + scheduler = _base_scheduler() + + rebuilt = policy.set_flow_shift(scheduler, 3.0) + assert rebuilt is not scheduler + assert isinstance(rebuilt, UniPCMultistepScheduler) + assert float(rebuilt.config.flow_shift) == 3.0 + + def test_current_shift_read_from_scheduler_config(self): + """No separate shift-tracking state: a second call with the same target + on the rebuilt instance is a no-op; restoring rebuilds again.""" + policy = _base_policy() + rebuilt = policy.set_flow_shift(_base_scheduler(), 3.0) + assert policy.set_flow_shift(rebuilt, 3.0) is rebuilt + + restored = policy.set_flow_shift(rebuilt, 1.0) + assert restored is not rebuilt + assert float(restored.config.flow_shift) == 1.0 + + def test_distilled_is_structural_noop(self): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + assert policy.set_flow_shift(scheduler, 3.0) is scheduler + + +class TestSetTimesteps: + def test_distilled_programs_fixed_sigmas(self): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + policy.set_timesteps(scheduler, num_inference_steps=4, device="cpu") + expected = [s * 1000.0 for s in DISTILLED_SIGMAS] + assert torch.allclose(scheduler.timesteps.float(), torch.tensor(expected), atol=1e-3) + + def test_base_programs_step_count(self): + policy = _base_policy() + scheduler = _base_scheduler() + policy.set_timesteps(scheduler, num_inference_steps=7, device="cpu") + assert len(scheduler.timesteps) == 7 + + +class TestStochasticStepDeterminism: + """The seeded generator must fully determine the SDE noise trajectory.""" + + def _run_steps(self, seed): + policy = _distilled_policy() + scheduler = FlowMatchEulerDiscreteScheduler.from_config(DISTILLED_SCHEDULER_CONFIG) + policy.set_timesteps(scheduler, num_inference_steps=4, device="cpu") + generator = torch.Generator().manual_seed(seed) + kwargs = policy.scheduler_step_kwargs(generator) + + latents = torch.zeros(1, 4, 1, 2, 2) + velocity = torch.full_like(latents, 0.5) + for t in scheduler.timesteps: + latents = scheduler.step(velocity, t, latents, return_dict=False, **kwargs)[0] + return latents + + def test_same_seed_reproduces_sde_trajectory(self): + assert torch.equal(self._run_steps(seed=123), self._run_steps(seed=123)) + + def test_different_seeds_diverge(self): + assert not torch.equal(self._run_steps(seed=123), self._run_steps(seed=456)) + + +class TestGenerationDefaults: + def test_distilled_defaults_report_checkpoint_truth(self): + params = _bare_pipeline(sampling=_distilled_policy()).default_generation_params + assert params["num_inference_steps"] == 4 + assert params["guidance_scale"] == DISTILLED_GUIDANCE_SCALE + assert params["height"] is None # mode-dependent, resolved in infer() + assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] + + def test_base_defaults_leave_mode_dependent_fields_unset(self): + params = _bare_pipeline().default_generation_params + for field in ("height", "width", "num_inference_steps", "guidance_scale"): + assert params[field] is None + assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] + assert params["max_sequence_length"] == COSMOS3_720P_PARAMS["max_sequence_length"] + + +class TestInferModeResolution: + def _captured_forward_kwargs(self, pipeline, req): + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + pipeline.infer(req) + return captured + + def test_video_unset_resolves_to_video_table(self): + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("video")) + assert got["height"] == COSMOS3_720P_PARAMS["height"] + assert got["width"] == COSMOS3_720P_PARAMS["width"] + assert got["num_inference_steps"] == COSMOS3_720P_PARAMS["num_inference_steps"] + assert got["guidance_scale"] == COSMOS3_720P_PARAMS["guidance_scale"] + + def test_t2i_unset_resolves_to_t2i_table(self): + got = self._captured_forward_kwargs(_bare_pipeline(), _fake_request("image")) + assert got["height"] == COSMOS3_T2I_PARAMS["height"] + assert got["width"] == COSMOS3_T2I_PARAMS["width"] + assert got["num_inference_steps"] == COSMOS3_T2I_PARAMS["num_inference_steps"] + assert got["guidance_scale"] == COSMOS3_T2I_PARAMS["guidance_scale"] + + def test_explicit_values_pass_through(self): + req = _fake_request("image", height=512, num_inference_steps=20) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["height"] == 512 + assert got["num_inference_steps"] == 20 + assert got["width"] == COSMOS3_T2I_PARAMS["width"] + + def test_distilled_merged_defaults_pass_through(self): + req = _fake_request("image", num_inference_steps=4, guidance_scale=1.0) + got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) + assert got["num_inference_steps"] == 4 + assert got["guidance_scale"] == DISTILLED_GUIDANCE_SCALE + assert got["height"] == COSMOS3_T2I_PARAMS["height"] + + +class TestPipelineSchedulerLoading: + def test_distilled_checkpoint_loads_flow_match(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + pipeline = _bare_pipeline() + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.sampling.is_distilled + assert pipeline.sampling.fixed_sigmas == DISTILLED_SIGMAS + + def test_base_checkpoint_loads_unipc(self, tmp_path): + _write_scheduler_config(tmp_path, UNIPC_SCHEDULER_CONFIG) + pipeline = _bare_pipeline() + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.scheduler, UniPCMultistepScheduler) + assert not pipeline.sampling.is_distilled + assert pipeline.sampling.checkpoint_flow_shift == 1.0 + + def test_audio_scheduler_is_separate_same_class_instance(self, tmp_path): + _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) + pipeline = _bare_pipeline(audio_gen=True) + + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=SKIP_NON_SCHEDULER + ) + + assert isinstance(pipeline.audio_scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.audio_scheduler is not pipeline.scheduler + + +class TestWarmupAndForwardValidation: + def test_warmup_steps_follow_distilled_schedule(self): + assert _bare_pipeline(sampling=_distilled_policy()).default_warmup_steps == 4 + + def test_warmup_steps_base_default(self): + assert _bare_pipeline().default_warmup_steps == 2 # BasePipeline default + + @pytest.mark.parametrize( + "policy_factory, expected_guidance", + [(_distilled_policy, DISTILLED_GUIDANCE_SCALE), (_base_policy, 6.0)], + ) + def test_warmup_guidance_uses_pipeline_defaults(self, policy_factory, expected_guidance): + pipeline = _bare_pipeline(sampling=policy_factory()) + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + + pipeline._run_warmup(height=720, width=1280, num_frames=9, steps=4) + + assert captured["guidance_scale"] == expected_guidance + + @pytest.mark.parametrize( + "bad_kwargs", + [ + {"num_inference_steps": 10, "guidance_scale": 1.0}, + {"num_inference_steps": 4, "guidance_scale": 3.5}, + ], + ) + def test_forward_rejects_explicit_mismatch(self, bad_kwargs): + pipeline = _bare_pipeline(sampling=_distilled_policy()) + with pytest.raises(ValueError, match="distilled"): + pipeline.forward(prompt="x", seed=0, use_guardrails=False, **bad_kwargs) + + def test_distilled_rejects_image_conditioning(self): + """Without per-step re-anchoring, the stochastic scheduler corrupts the + conditioned frame; the request must fail rather than degrade silently.""" + pipeline = _bare_pipeline(sampling=_distilled_policy()) + with pytest.raises(ValueError, match="re-anchor"): + pipeline.forward( + prompt="x", + seed=0, + use_guardrails=False, + image="frame.png", + num_inference_steps=4, + guidance_scale=1.0, + ) + + def test_base_still_accepts_image_conditioning_path(self): + """Base checkpoints must not be caught by the distilled image rejection: + forward proceeds past validation (fails later on the bare test double).""" + pipeline = _bare_pipeline(sampling=_base_policy()) + with pytest.raises(AttributeError): + pipeline.forward(prompt="x", seed=0, use_guardrails=False, image="frame.png") + + @pytest.mark.parametrize("bad_output_type", ["imgae", "png", "", "both"]) + def test_invalid_output_type_raises(self, bad_output_type): + pipeline = _bare_pipeline() + with pytest.raises(ValueError, match="output_type"): + pipeline.forward(prompt="x", seed=0, use_guardrails=False, output_type=bad_output_type) + + +class _AdditiveScheduler: + """step(v, t, x) = x + v, so final latents are the exact sum of all + predictions. The strict signature also pins that the loop passes no silent + extra step kwargs unless the caller supplies them.""" + + def __init__(self, timesteps): + self.timesteps = timesteps + + def step(self, model_output, timestep, sample, return_dict=False): + assert return_dict is False + return (sample + model_output,) + + +class _GeneratorRecordingScheduler(_AdditiveScheduler): + def __init__(self, timesteps): + super().__init__(timesteps) + self.generators = [] + + def step(self, model_output, timestep, sample, return_dict=False, generator=None): + self.generators.append(generator) + return super().step(model_output, timestep, sample, return_dict) + + +def _denoise_ready_pipeline() -> Cosmos3OmniMoTPipeline: + return _bare_pipeline( + pipeline_config=SimpleNamespace(visual_gen_mapping=None), + cache_accelerator=None, + _predenoise_pending=False, + _postdenoise_pending=False, + _is_warmup=False, + _profile_range=None, + ) + + +class TestDistilledDenoiseLoop: + POS_IDS = torch.arange(8).unsqueeze(0) + NEG_IDS = torch.arange(8).unsqueeze(0) + 100 + POS_MASK = torch.ones(1, 8, dtype=torch.long) + NEG_MASK = torch.zeros(1, 8, dtype=torch.long) + + def _run(self, scheduler=None, scheduler_step_kwargs=None): + pipeline = _denoise_ready_pipeline() + timesteps = torch.tensor([s * 1000.0 for s in DISTILLED_SIGMAS]) + scheduler = scheduler if scheduler is not None else _AdditiveScheduler(timesteps) + calls = [] + + def forward_fn(latent_input, extra_streams, step_index, timestep, embeds, extras): + calls.append( + { + "batch": latent_input.shape[0], + "timestep": float(timestep[0]), + "text_ids": extras["text_ids"], + } + ) + return torch.full_like(latent_input, 0.5) + + latents = torch.zeros(1, 4, 3, 2, 2) + result = pipeline.denoise( + latents=latents, + scheduler=scheduler, + prompt_embeds=self.POS_IDS, + neg_prompt_embeds=self.NEG_IDS, + guidance_scale=DISTILLED_GUIDANCE_SCALE, + forward_fn=forward_fn, + extra_cfg_tensors={ + "text_ids": (self.POS_IDS, self.NEG_IDS), + "text_mask": (self.POS_MASK, self.NEG_MASK), + }, + scheduler_step_kwargs=scheduler_step_kwargs, + ) + return result, calls + + def test_guidance_one_single_forward_per_step(self): + result, calls = self._run() + + assert len(calls) == 4, "one forward per distilled step, no CFG branch" + assert all(c["batch"] == 1 for c in calls), "no CFG batch duplication" + assert all(c["text_ids"] is self.POS_IDS for c in calls), "positive prompt only" + assert [c["timestep"] for c in calls] == pytest.approx( + [s * 1000.0 for s in DISTILLED_SIGMAS], abs=1e-3 + ) + assert torch.all(result == 2.0) # 4 additive steps of +0.5 from 0 + + def test_scheduler_step_kwargs_reach_every_step(self): + generator = torch.Generator().manual_seed(7) + timesteps = torch.tensor([s * 1000.0 for s in DISTILLED_SIGMAS]) + scheduler = _GeneratorRecordingScheduler(timesteps) + + result, calls = self._run( + scheduler=scheduler, + scheduler_step_kwargs=_distilled_policy().scheduler_step_kwargs(generator), + ) + + assert len(calls) == 4 + assert scheduler.generators == [generator] * 4 + assert torch.all(result == 2.0) + + +class TestRegistryDispatch: + def test_model_index_class_name_dispatches(self, tmp_path): + with open(tmp_path / "model_index.json", "w") as f: + json.dump({"_class_name": "Cosmos3OmniPipeline"}, f) + assert AutoPipeline._detect_from_checkpoint(str(tmp_path)) == "Cosmos3OmniMoTPipeline" + + def test_hf_id_registered(self): + entry = PIPELINE_REGISTRY["Cosmos3OmniMoTPipeline"] + assert "nvidia/Cosmos3-Super-Text2Image-4Step" in entry.hf_ids diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index c0329b71a060..d1993549615b 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -30,7 +30,6 @@ from pathlib import Path os.environ["TLLM_DISABLE_MPI"] = "1" -os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" import PIL.Image import pytest @@ -46,7 +45,7 @@ from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs -pytestmark = pytest.mark.cosmos3 +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] @pytest.fixture(autouse=True, scope="module") diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 76e6c1a01a8c..3ebcb08290da 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -20,20 +20,24 @@ import gc import os from pathlib import Path +from types import SimpleNamespace os.environ["TLLM_DISABLE_MPI"] = "1" -os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" import pytest import torch from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig -from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + PRETRAINED_CONFIG_COMPAT_DEFAULTS, + Cosmos3VFMTransformer, + apply_pretrained_config_compat_defaults, +) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs -pytestmark = pytest.mark.cosmos3 +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] @pytest.fixture(autouse=True, scope="module") @@ -468,3 +472,36 @@ def test_load_fp8_quantization(self, quant_algo: str): del pipeline gc.collect() torch.cuda.empty_cache() + + +# --- CPU-only coverage: checkpoint config schema compatibility --- + + +class TestConfigCompatDefaults: + """Newer diffusers conversions omit fields older ones carried explicitly.""" + + def test_new_schema_gets_defaults(self): + config = SimpleNamespace(hidden_size=64, rope_axes_dim=[4, 2, 2]) + apply_pretrained_config_compat_defaults(config) + for key, value in PRETRAINED_CONFIG_COMPAT_DEFAULTS.items(): + assert getattr(config, key) == value + + def test_old_schema_untouched(self): + # Every field deliberately differs from its compat default, so an + # overwrite of any one of them fails its assertion. + config = SimpleNamespace( + position_embedding_type="rope_3d", + max_position_embeddings=12345, + temporal_compression_factor_sound=7, + ) + apply_pretrained_config_compat_defaults(config) + assert config.position_embedding_type == "rope_3d" + assert config.max_position_embeddings == 12345 + assert config.temporal_compression_factor_sound == 7 + + def test_idempotent(self): + config = SimpleNamespace(hidden_size=64) + apply_pretrained_config_compat_defaults(config) + snapshot = vars(config).copy() + apply_pretrained_config_compat_defaults(config) + assert vars(config) == snapshot