From 1a09ad7a2e5391318a5bbb608fb82ac6071fb12e Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:12:28 +0800 Subject: [PATCH 1/3] [None][refactor] split visual gen pipeline and model configs Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/__init__.py | 3 +- tensorrt_llm/_torch/visual_gen/config.py | 168 +++++++++++++++--- .../models/cosmos3/pipeline_cosmos3.py | 6 +- .../models/cosmos3/transformer_cosmos3.py | 6 +- .../visual_gen/models/flux/pipeline_flux.py | 10 +- .../visual_gen/models/flux/pipeline_flux2.py | 10 +- .../models/flux/transformer_flux.py | 6 +- .../models/flux/transformer_flux2.py | 6 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 7 +- .../models/ltx2/transformer_ltx2.py | 9 +- .../_torch/visual_gen/models/modeling.py | 29 +++ .../models/qwen_image/pipeline_qwen_image.py | 17 +- .../qwen_image/transformer_qwen_image.py | 7 +- .../visual_gen/models/wan/pipeline_wan.py | 17 +- .../visual_gen/models/wan/pipeline_wan_i2v.py | 16 +- .../visual_gen/models/wan/transformer_wan.py | 7 +- tensorrt_llm/_torch/visual_gen/pipeline.py | 19 +- .../_torch/visual_gen/pipeline_loader.py | 16 +- .../_torch/visual_gen/pipeline_registry.py | 10 +- tensorrt_llm/visual_gen/args.py | 2 +- .../_torch/visual_gen/test_ltx2_pipeline.py | 18 +- .../_torch/visual_gen/test_model_loader.py | 14 +- .../test_qwen_image_pipeline_config.py | 4 +- .../_torch/visual_gen/test_visual_gen_args.py | 6 +- .../_torch/visual_gen/test_wan_transformer.py | 10 +- 25 files changed, 294 insertions(+), 129 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/models/modeling.py diff --git a/tensorrt_llm/_torch/visual_gen/__init__.py b/tensorrt_llm/_torch/visual_gen/__init__.py index 2be964b60215..f7c5ca753ccd 100644 --- a/tensorrt_llm/_torch/visual_gen/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/__init__.py @@ -12,7 +12,7 @@ from tensorrt_llm._torch.visual_gen.output import PipelineOutput from .checkpoints import WeightLoader -from .config import DiffusionModelConfig +from .config import DiffusionModelConfig, DiffusionPipelineConfig from .mapping import VisualGenMapping from .models import AutoPipeline, BasePipeline, WanPipeline from .pipeline_loader import PipelineLoader @@ -20,6 +20,7 @@ __all__ = [ "DiffusionModelConfig", + "DiffusionPipelineConfig", "PipelineComponent", "WeightLoader", "PipelineLoader", diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 34c7b6665621..d709fab278e6 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -12,7 +12,7 @@ # 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. -"""Internal DiffusionModelConfig and loading helpers.""" +"""Internal VisualGen pipeline and model configuration helpers.""" import json from pathlib import Path @@ -85,25 +85,84 @@ def create_attention_metadata_state() -> Dict[str, Any]: return {"metadata": None, "capacity": (0, 0)} +class _VisualGenConfigBase(BaseModel): + """Base for internal VisualGen configs that carry runtime objects.""" + + # Pydantic reserves `model_config` for class-level settings. This is not + # a VisualGen model config; it lets fields hold objects such as Mapping. + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class DiffusionModelConfig(_VisualGenConfigBase): + """Internal config for one TRT-LLM VisualGen model component.""" + + component_name: Optional[str] = None + pretrained_config: Optional[Any] = None + mapping: Mapping = PydanticField(default_factory=Mapping) + skip_create_weights_in_init: bool = False + force_dynamic_quantization: bool = False + allreduce_strategy: AllReduceStrategy = PydanticField(default=AllReduceStrategy.NCCL) + extra_attrs: Dict = PydanticField(default_factory=dict) + + # Unified parallelism mapping copied from the owning pipeline config. + visual_gen_mapping: Optional[Any] = None # VisualGenMapping (lazy import) + + dynamic_weight_quant: bool = False + + # Shared runtime configs copied from the owning pipeline config. + quant_config: QuantConfig = PydanticField(default_factory=QuantConfig) + # Per-layer quant (from load_diffusion_quant_config layer_quant_config; None until mixed-precision parsing exists) + quant_config_dict: Optional[Dict[str, QuantConfig]] = None + compilation: CompilationConfig = PydanticField(default_factory=CompilationConfig) + torch_compile: TorchCompileConfig = PydanticField(default_factory=TorchCompileConfig) + cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) + attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) + attention_metadata_state: Optional[Dict[str, Any]] = None + parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) + cache: Optional[CacheConfig] = None + + # Observability — flat field mirrors VisualGenArgs.enable_layerwise_nvtx_marker. + enable_layerwise_nvtx_marker: bool = False + + @property + def cache_backend(self) -> Optional[CacheBackendName]: + return self.cache.cache_backend if self.cache is not None else None # type: ignore[return-value] + + @property + def teacache(self) -> Optional[TeaCacheConfig]: + return self.cache if isinstance(self.cache, TeaCacheConfig) else None + + @property + def cache_dit(self) -> Optional[CacheDiTConfig]: + return self.cache if isinstance(self.cache, CacheDiTConfig) else None + + @property + def torch_dtype(self) -> "torch.dtype": + """Get the torch dtype of the model (default: bfloat16).""" + return torch.bfloat16 + + def get_quant_config(self, name: Optional[str] = None) -> QuantConfig: + """Get quantization config for a layer or global. Resembles LLM ModelConfig.get_quant_config.""" + if name is None or self.quant_config_dict is None: + return self.quant_config + if name in self.quant_config_dict: + return self.quant_config_dict[name] + return self.quant_config + + # ============================================================================= -# DiffusionModelConfig - Internal configuration (merged/parsed) +# DiffusionPipelineConfig - Internal pipeline configuration (merged/parsed) # ============================================================================= -class DiffusionModelConfig(BaseModel): - """Internal ModelConfig for diffusion models. +class DiffusionPipelineConfig(_VisualGenConfigBase): + """Internal config for an entire VisualGen pipeline. - This is created by PipelineLoader from VisualGenArgs + checkpoint. - Contains merged/parsed config from: - - pretrained_config: From checkpoint/config.json - - quant_config: From checkpoint or user quant config - - Sub-configs: From VisualGenArgs (pipeline, attention, teacache) - - visual_gen_mapping: Populated by setup_visual_gen_mapping() from ParallelConfig + This is created by PipelineLoader from VisualGenArgs + checkpoint and owns + pipeline/runtime state plus one DiffusionModelConfig per model component. """ - model_config = ConfigDict(arbitrary_types_allowed=True) - - pretrained_config: Optional[Any] = None + model_configs: Dict[str, DiffusionModelConfig] = PydanticField(default_factory=dict) mapping: Mapping = PydanticField(default_factory=Mapping) skip_create_weights_in_init: bool = False force_dynamic_quantization: bool = False @@ -127,15 +186,17 @@ class DiffusionModelConfig(BaseModel): parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None - # Merged per-family pipeline_config: registry-entry defaults overlaid - # with the user-supplied VisualGenArgs.pipeline_config dict (user - # values win). Validated against the registry entry's `defaults` - # before assignment, so unknown keys never reach here. - pipeline_config: Dict[str, Any] = PydanticField(default_factory=dict) - # Observability — flat field mirrors VisualGenArgs.enable_layerwise_nvtx_marker. enable_layerwise_nvtx_marker: bool = False + @property + def primary_model_config(self) -> DiffusionModelConfig: + return self.model_configs["transformer"] + + @property + def primary_pretrained_config(self) -> Any: + return self.primary_model_config.pretrained_config + @property def cache_backend(self) -> Optional[CacheBackendName]: return self.cache.cache_backend if self.cache is not None else None # type: ignore[return-value] @@ -161,6 +222,33 @@ def get_quant_config(self, name: Optional[str] = None) -> QuantConfig: return self.quant_config_dict[name] return self.quant_config + def _make_model_config( + self, + component_name: str, + model_pretrained_config: Any, + ) -> DiffusionModelConfig: + return DiffusionModelConfig( + component_name=component_name, + pretrained_config=model_pretrained_config, + mapping=self.mapping, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + allreduce_strategy=self.allreduce_strategy, + extra_attrs=self.extra_attrs, + visual_gen_mapping=self.visual_gen_mapping, + dynamic_weight_quant=self.dynamic_weight_quant, + quant_config=self.quant_config, + quant_config_dict=self.quant_config_dict, + compilation=self.compilation, + torch_compile=self.torch_compile, + cuda_graph=self.cuda_graph, + attention=self.attention, + attention_metadata_state=self.attention_metadata_state, + parallel=self.parallel, + cache=self.cache, + enable_layerwise_nvtx_marker=self.enable_layerwise_nvtx_marker, + ) + @staticmethod def load_diffusion_quant_config( quant_config_dict: dict, @@ -346,12 +434,12 @@ def from_pretrained( checkpoint_dir: str, args: Optional["VisualGenArgs"] = None, **kwargs, - ) -> "DiffusionModelConfig": + ) -> "DiffusionPipelineConfig": """ Load config from pretrained checkpoint. Called by PipelineLoader with VisualGenArgs: - config = DiffusionModelConfig.from_pretrained( + config = DiffusionPipelineConfig.from_pretrained( checkpoint_dir=args.model, args=args, ) @@ -404,6 +492,7 @@ def from_pretrained( # Discover pipeline components (diffusers layout) components = discover_pipeline_components(checkpoint_path) + component_config_dicts: Dict[str, Dict[str, Any]] = {} if components: # ---------- Diffusers directory layout ---------- @@ -415,8 +504,11 @@ def from_pretrained( if not config_path.exists(): raise ValueError(f"Config not found at {config_path}") - with open(config_path) as f: - config_dict = json.load(f) + for component_name, component_config_path in components.items(): + with open(component_config_path) as f: + component_config_dicts[component_name] = json.load(f) + + config_dict = component_config_dicts[component] pretrained_config = SimpleNamespace(**config_dict) # Ensure _name_or_path is set so TeaCache coefficient matching works. @@ -439,6 +531,10 @@ def from_pretrained( if native_config is not None: transformer_dict = native_config.get("transformer", {}) + component_config_dicts["transformer"] = transformer_dict + transformer_2_dict = native_config.get("transformer_2") + if isinstance(transformer_2_dict, dict): + component_config_dicts["transformer_2"] = transformer_2_dict pretrained_config = SimpleNamespace(**transformer_dict) if not getattr(pretrained_config, "_name_or_path", None): pretrained_config._name_or_path = str(checkpoint_path) @@ -551,8 +647,7 @@ def from_pretrained( create_attention_metadata_state() if attention_cfg.backend == "TRTLLM" else None ) - return cls( - pretrained_config=pretrained_config, + pipeline_config = cls( quant_config=quant_config, quant_config_dict=quant_config_dict, dynamic_weight_quant=dynamic_weight_quant, @@ -566,8 +661,29 @@ def from_pretrained( parallel=parallel_cfg, cache=cache_cfg, enable_layerwise_nvtx_marker=enable_layerwise_nvtx_marker, - pipeline_config=resolved_pipeline_config, skip_create_weights_in_init=True, extra_attrs=extra_attrs, **kwargs, ) + + for component_name, config_dict in component_config_dicts.items(): + if component_name == component: + component_pretrained_config = pretrained_config + else: + component_pretrained_config = SimpleNamespace(**config_dict) + if not getattr(component_pretrained_config, "_name_or_path", None): + component_pretrained_config._name_or_path = getattr( + pretrained_config, "_name_or_path", "" + ) + pipeline_config.model_configs[component_name] = pipeline_config._make_model_config( + component_name, + component_pretrained_config, + ) + + if not pipeline_config.model_configs: + pipeline_config.model_configs["transformer"] = pipeline_config._make_model_config( + "transformer", + pretrained_config, + ) + + return pipeline_config 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 d01e71d998ec..06b6f7fdad66 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -63,12 +63,12 @@ doc="Cosmos3 Omnimodal world models.", ) class Cosmos3OmniMoTPipeline(BasePipeline): - def __init__(self, model_config): - super().__init__(model_config) + def __init__(self, pipeline_config): + super().__init__(pipeline_config) def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") - self.transformer = Cosmos3VFMTransformer(self.model_config) + self.transformer = Cosmos3VFMTransformer(self.model_configs["transformer"]) def load_weights(self, weights: dict) -> None: if self.transformer is not None and hasattr(self.transformer, "load_weights"): 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 10de6fc87235..82e8d7fa65b6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder @@ -641,10 +642,9 @@ def forward( return cached_kv -class Cosmos3VFMTransformer(nn.Module): +class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): - super().__init__() - self.model_config = model_config + super().__init__(model_config) pretrained_config = model_config.pretrained_config self.hidden_size = pretrained_config.hidden_size diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index c43345f9c8c4..883c8b3269e6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -50,16 +50,16 @@ class FluxPipeline(BasePipeline): Supports FLUX.1-dev (50 steps, guidance) and FLUX.1-schnell (4 steps, no guidance). """ - def __init__(self, model_config): + def __init__(self, pipeline_config): if ( - model_config.visual_gen_mapping is not None - and model_config.visual_gen_mapping.cfg_size != 1 + pipeline_config.visual_gen_mapping is not None + and pipeline_config.visual_gen_mapping.cfg_size != 1 ): raise ValueError( "FluxPipeline does not support CFG parallelism. Please set cfg_size to 1." ) - super().__init__(model_config) + super().__init__(pipeline_config) @staticmethod def _compute_flux_timestep_embedding( @@ -121,7 +121,7 @@ def warmup_cache_key(self, height: int, width: int, **kwargs) -> tuple: def _init_transformer(self) -> None: """Initialize FLUX transformer with quantization support.""" logger.info("Creating FLUX transformer with quantization support...") - self.transformer = FluxTransformer2DModel(model_config=self.model_config) + self.transformer = FluxTransformer2DModel(model_config=self.model_configs["transformer"]) def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 8675a0387e79..81a40a50066a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -119,16 +119,16 @@ class Flux2Pipeline(BasePipeline): # Default for backward compatibility (FLUX.2-dev) HIDDEN_STATE_LAYERS: Tuple[int, ...] = (10, 20, 30) - def __init__(self, model_config): + def __init__(self, pipeline_config): if ( - model_config.visual_gen_mapping is not None - and model_config.visual_gen_mapping.cfg_size != 1 + pipeline_config.visual_gen_mapping is not None + and pipeline_config.visual_gen_mapping.cfg_size != 1 ): raise ValueError( "Flux2Pipeline does not support CFG parallelism. Please set cfg_size to 1." ) - super().__init__(model_config) + super().__init__(pipeline_config) @staticmethod def _compute_flux2_timestep_embedding( @@ -189,7 +189,7 @@ def warmup_cache_key(self, height: int, width: int, **kwargs) -> tuple: def _init_transformer(self) -> None: """Initialize FLUX.2 transformer with quantization support.""" logger.info("Creating FLUX.2 transformer with quantization support...") - self.transformer = Flux2Transformer2DModel(model_config=self.model_config) + self.transformer = Flux2Transformer2DModel(model_config=self.model_configs["transformer"]) def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py index e87ace476d74..99852d3f6f6b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py @@ -34,6 +34,7 @@ from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import FluxJointAttnMLPProj from tensorrt_llm._torch.visual_gen.models.flux.pos_embed_flux import FluxPosEmbed +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig @@ -554,7 +555,7 @@ def forward( return encoder_hidden_states, hidden_states -class FluxTransformer2DModel(nn.Module): +class FluxTransformer2DModel(BaseDiffusionModel): """FLUX Transformer model for text-to-image generation. This is the native TRT-LLM implementation of FLUX transformer. @@ -572,8 +573,7 @@ class FluxTransformer2DModel(nn.Module): """ def __init__(self, model_config: DiffusionModelConfig): - super().__init__() - self.model_config = model_config + super().__init__(model_config) vgm = model_config.visual_gen_mapping num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 24) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py index 15dce09d9565..0fb2d53311df 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py @@ -41,6 +41,7 @@ AdaLayerNormContinuous, _remap_checkpoint_keys, ) +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig @@ -417,7 +418,7 @@ def forward( # ============================================================================= -class Flux2Transformer2DModel(nn.Module): +class Flux2Transformer2DModel(BaseDiffusionModel): """FLUX.2 Transformer model for image generation (Native TRT-LLM). This implements the full FLUX.2 architecture matching HuggingFace diffusers: @@ -433,8 +434,7 @@ def __init__(self, model_config: DiffusionModelConfig): Args: model_config: DiffusionModelConfig instance (from DiffusionModelLoader) """ - super().__init__() - self.model_config = model_config + super().__init__(model_config) vgm = model_config.visual_gen_mapping num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 48) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index c66d795bc47f..ed25a96ed936 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -629,7 +629,7 @@ def resolve_variant(cls, config): logger.info(f"{LTX2_FORCE_ONE_STAGE_ENV} is enabled; forcing one-stage LTX2 pipeline.") return cls - checkpoint_path = getattr(config.pretrained_config, "_name_or_path", "") + checkpoint_path = getattr(config.primary_pretrained_config, "_name_or_path", "") if checkpoint_path: config.extra_attrs.update( resolve_ltx2_pipeline_extra_attrs(Path(checkpoint_path), config.extra_attrs) @@ -739,7 +739,8 @@ def _init_transformer(self) -> None: "Quantized attention is not yet supported for the LTX-2 pipeline." ) - cfg = self.model_config.pretrained_config + model_config = self.model_configs["transformer"] + cfg = model_config.pretrained_config rope_type = LTXRopeType(getattr(cfg, "rope_type", "interleaved")) freq_prec = getattr(cfg, "frequencies_precision", False) @@ -780,7 +781,7 @@ def _init_transformer(self) -> None: rope_type=rope_type, double_precision_rope=double_precision_rope, apply_gated_attention=apply_gated_attention, - model_config=self.model_config, + model_config=model_config, ) self.transformer._transformer_config = vars(cfg) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index ce09797c389c..1631573f28f3 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -30,6 +30,7 @@ from tensorrt_llm._torch.modules.linear import Linear, WeightMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder @@ -928,7 +929,7 @@ def is_audio_enabled(self) -> bool: return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) -class LTXModel(nn.Module): +class LTXModel(BaseDiffusionModel): """LTX-2 transformer built from TRT-LLM primitives. Native implementation using optimized TRT-LLM Linear, RMSNorm, MLP, and @@ -966,8 +967,10 @@ def __init__( apply_gated_attention: bool = False, model_config: Optional["DiffusionModelConfig"] = None, ): - super().__init__() - self.model_config = model_config + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + + model_config = model_config or DiffusionModelConfig() + super().__init__(model_config) self.model_type = model_type self.use_middle_indices_grid = use_middle_indices_grid self.rope_type = rope_type diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py new file mode 100644 index 000000000000..899feffb429e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -0,0 +1,29 @@ +# 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. +"""Base classes for VisualGen model components.""" + +import torch.nn as nn + +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + + +class BaseDiffusionModel(nn.Module): + """Base class for TRT-LLM VisualGen model components.""" + + def __init__(self, model_config: DiffusionModelConfig): + super().__init__() + self.model_config = model_config + self.component_name = model_config.component_name + self.pretrained_config = model_config.pretrained_config diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py index 1c42f96d458b..aed7e29554e9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py @@ -76,8 +76,8 @@ class QwenImagePipeline(BasePipeline): # either version. DEFAULT_GENERATION_PARAMS = _DEFAULT_GENERATION_PARAMS - def __init__(self, model_config): - super().__init__(model_config) + def __init__(self, pipeline_config): + super().__init__(pipeline_config) # Qwen-Image uses 8x VAE downsample + 2x2 patch packing. Both # scheduler and image-prep assume a latent grid divisible by # (vae_scale_factor * 2 == 16). vae_scale_factor is updated by @@ -122,12 +122,11 @@ def resolution_multiple_of(self) -> Tuple[int, int]: # ------------------------------------------------------------------ def _init_transformer(self) -> None: logger.info("Creating Qwen-Image transformer") - # ``pretrained_config`` on the DiffusionModelConfig is populated - # from ``/transformer/config.json`` as a SimpleNamespace by - # ``DiffusionModelConfig.from_pretrained``. Read the fields we - # care about with sensible defaults (matching the Qwen-Image 20B - # reference model). - pretrained = getattr(self.model_config, "pretrained_config", None) + model_config = self.model_configs["transformer"] + # ``pretrained_config`` is populated from + # ``/transformer/config.json``. Read the fields we care + # about with defaults matching the Qwen-Image 20B reference model. + pretrained = getattr(model_config, "pretrained_config", None) def _cfg(name: str, default): if pretrained is None: @@ -137,7 +136,7 @@ def _cfg(name: str, default): return getattr(pretrained, name, default) self.transformer = QwenImageTransformer2DModel( - model_config=self.model_config, + model_config=model_config, patch_size=_cfg("patch_size", 2), in_channels=_cfg("in_channels", 64), out_channels=_cfg("out_channels", 16), diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py index 9b3330d51011..f06db14a45b6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py @@ -29,6 +29,7 @@ from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader @@ -734,7 +735,7 @@ def forward( # =========================================================================== -class QwenImageTransformer2DModel(nn.Module): +class QwenImageTransformer2DModel(BaseDiffusionModel): """Qwen-Image 20B MMDiT transformer. Mirrors ``diffusers.models.transformers.transformer_qwenimage.QwenImageTransformer2DModel`` @@ -756,8 +757,8 @@ def __init__( axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), attn_backend: str = "sdpa", ): - super().__init__() - self.model_config = model_config or DiffusionModelConfig() + model_config = model_config or DiffusionModelConfig() + super().__init__(model_config) self.attn_backend = attn_backend self.patch_size = patch_size diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index f314f33fad3f..c53dea7f0e13 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -89,23 +89,24 @@ doc="Wan 2.1 & 2.2 text-to-video family.", ) class WanPipeline(BasePipeline): - def __init__(self, model_config): + def __init__(self, pipeline_config): # Wan2.2 A14B two-stage denoising parameters self.transformer_2 = None - self.boundary_ratio = getattr(model_config.pretrained_config, "boundary_ratio", None) - self.expand_timesteps = getattr(model_config.pretrained_config, "expand_timesteps", False) + primary_pretrained_config = pipeline_config.primary_pretrained_config + self.boundary_ratio = getattr(primary_pretrained_config, "boundary_ratio", None) + self.expand_timesteps = getattr(primary_pretrained_config, "expand_timesteps", False) # Derived model type flags self.is_wan22_14b = self.boundary_ratio is not None self.is_wan22_5b = self.expand_timesteps # Validate TeaCache compatibility before allocating GPU memory - if (self.is_wan22_14b or self.is_wan22_5b) and model_config.cache_backend == "teacache": + if (self.is_wan22_14b or self.is_wan22_5b) and pipeline_config.cache_backend == "teacache": raise ValueError( "TeaCache is not supported for Wan 2.2 models. " "Use cache_backend='none' or 'cache_dit' (not 'teacache')." ) - super().__init__(model_config) + super().__init__(pipeline_config) def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): """Compute timestep embedding for WAN transformer. @@ -176,12 +177,14 @@ def resolution_multiple_of(self): def _init_transformer(self) -> None: logger.info("Creating WAN transformer with quantization support...") - self.transformer = WanTransformer3DModel(model_config=self.model_config) + self.transformer = WanTransformer3DModel(model_config=self.model_configs["transformer"]) # Wan2.2 A14B: create second transformer for two-stage denoising if self.is_wan22_14b: logger.info("Creating second transformer for Wan2.2 A14B two-stage denoising...") - self.transformer_2 = WanTransformer3DModel(model_config=self.model_config) + self.transformer_2 = WanTransformer3DModel( + model_config=self.model_configs["transformer_2"] + ) def load_standard_components( self, diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index ca9e53ecf9a8..dd70e3005612 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -88,20 +88,22 @@ doc="Wan 2.1 & 2.2 image-to-video family.", ) class WanImageToVideoPipeline(BasePipeline): - def __init__(self, model_config): + def __init__(self, pipeline_config): # Wan2.2 14B two-stage denoising parameters self.transformer_2 = None - self.boundary_ratio = getattr(model_config.pretrained_config, "boundary_ratio", None) + self.boundary_ratio = getattr( + pipeline_config.primary_pretrained_config, "boundary_ratio", None + ) self.is_wan22_14b = self.boundary_ratio is not None # Validate TeaCache compatibility before allocating GPU memory - if self.is_wan22_14b and model_config.cache_backend == "teacache": + if self.is_wan22_14b and pipeline_config.cache_backend == "teacache": raise ValueError( "TeaCache is not supported for Wan 2.2 models. " "Use cache_backend='none' or 'cache_dit' (not 'teacache')." ) - super().__init__(model_config) + super().__init__(pipeline_config) def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): """Compute timestep embedding for Wan I2V transformer. @@ -166,12 +168,14 @@ def resolution_multiple_of(self): def _init_transformer(self) -> None: logger.info("Creating WAN I2V transformer with quantization support...") - self.transformer = WanTransformer3DModel(model_config=self.model_config) + self.transformer = WanTransformer3DModel(model_config=self.model_configs["transformer"]) # Wan2.2: Optionally create second transformer for two-stage denoising if self.boundary_ratio is not None: logger.info("Creating second transformer for Wan2.2 I2V two-stage denoising...") - self.transformer_2 = WanTransformer3DModel(model_config=self.model_config) + self.transformer_2 = WanTransformer3DModel( + model_config=self.model_configs["transformer_2"] + ) def load_standard_components( self, diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index 0bac474df0ff..084c03754728 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -12,6 +12,7 @@ from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.modules.rms_norm import RMSNormTPAware from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader @@ -474,16 +475,14 @@ def forward( return x -class WanTransformer3DModel(nn.Module): +class WanTransformer3DModel(BaseDiffusionModel): _supports_gradient_checkpointing = True def __init__( self, model_config: DiffusionModelConfig, ): - super().__init__() - - self.model_config = model_config + super().__init__(model_config) vgm = model_config.visual_gen_mapping diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 9d19d7cf778e..10e87d3fe43d 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -98,7 +98,7 @@ def _parse_profile_range(): if TYPE_CHECKING: from .cache import CacheAccelerator - from .config import DiffusionModelConfig + from .config import DiffusionPipelineConfig class BasePipeline(nn.Module): @@ -107,7 +107,7 @@ class BasePipeline(nn.Module): """ @classmethod - def resolve_variant(cls, config: "DiffusionModelConfig") -> Type["BasePipeline"]: + def resolve_variant(cls, config: "DiffusionPipelineConfig") -> Type["BasePipeline"]: """Return *cls* or a more specialized subclass based on *config*. Override in subclasses to select a variant pipeline at creation @@ -117,11 +117,13 @@ def resolve_variant(cls, config: "DiffusionModelConfig") -> Type["BasePipeline"] """ return cls - def __init__(self, model_config: "DiffusionModelConfig"): + def __init__(self, pipeline_config: "DiffusionPipelineConfig"): super().__init__() - self.model_config = model_config - self.config = model_config.pretrained_config - self.mapping: Mapping = getattr(model_config, "mapping", None) or Mapping() + self.pipeline_config = pipeline_config + self.model_config = pipeline_config + self.model_configs = pipeline_config.model_configs + self.config = pipeline_config.primary_pretrained_config + self.mapping: Mapping = getattr(pipeline_config, "mapping", None) or Mapping() self._cuda_graph_runners: Dict[str, CUDAGraphRunner] = {} self._parallel_vae_enabled: bool = False self._warmed_up_shapes: Set[tuple] = set() @@ -401,10 +403,7 @@ def _apply_teacache_coefficients(self, coefficients: Optional[Dict]) -> None: if not coefficients: return teacache_cfg = self.model_config.teacache - checkpoint_path = ( - getattr(getattr(self.model_config, "pretrained_config", None), "_name_or_path", "") - or "" - ) + checkpoint_path = getattr(self.model_config.primary_pretrained_config, "_name_or_path", "") matched = False for model_size, coeff_data in coefficients.items(): if model_size.lower() in checkpoint_path.lower(): diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index ac61ce729fbc..d4e0a7f1295e 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -2,7 +2,7 @@ Model loader for diffusion pipelines. Flow: -1. Load config via DiffusionModelConfig.from_pretrained() +1. Load config via DiffusionPipelineConfig.from_pretrained() 2. Create pipeline via AutoPipeline.from_config() with MetaInit 3. Load weights with on-the-fly quantization if dynamic_weight_quant=True 4. Call pipeline.post_load_weights() @@ -28,7 +28,7 @@ from tensorrt_llm.visual_gen.args import VisualGenArgs from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxConfig, apply_skip_softmax_overrides -from .config import DiffusionModelConfig +from .config import DiffusionPipelineConfig from .mapping import VisualGenMapping from .models import AutoPipeline from .pipeline_registry import PIPELINE_REGISTRY, PipelineComponent @@ -140,7 +140,7 @@ def _resolve_pipeline_config(self, checkpoint_dir: str) -> dict: ) return {**entry.defaults, **user_pipeline_config} - def _setup_visual_gen_mapping(self, config: DiffusionModelConfig) -> None: + def _setup_visual_gen_mapping(self, config: DiffusionPipelineConfig) -> None: ws = dist.get_world_size() if dist.is_initialized() else 1 rk = dist.get_rank() if dist.is_initialized() else 0 attn2d_row, attn2d_col = self.args.parallel_config.attn2d_size @@ -155,8 +155,12 @@ def _setup_visual_gen_mapping(self, config: DiffusionModelConfig) -> None: tp_size=self.args.parallel_config.tp_size, parallel_vae_size=self.args.parallel_config.parallel_vae_size, ) + llm_mapping = vgm.to_llm_mapping() config.visual_gen_mapping = vgm - config.mapping = vgm.to_llm_mapping() + config.mapping = llm_mapping + for model_config in config.model_configs.values(): + model_config.visual_gen_mapping = vgm + model_config.mapping = llm_mapping def load( self, @@ -169,7 +173,7 @@ def load( Flow: 1. Resolve checkpoint_dir (local path or HuggingFace Hub model ID) - 2. Load config via DiffusionModelConfig.from_pretrained() + 2. Load config via DiffusionPipelineConfig.from_pretrained() 3. Create pipeline via AutoPipeline.from_config() with MetaInit 4. Load transformer weights via pipeline.load_transformer_weights() 5. Load auxiliary components (VAE, text_encoder) @@ -207,7 +211,7 @@ def load( # Merge pretrained checkpoint config with user-provided VisualGenArgs # ===================================================================== logger.info(f"Loading config from {checkpoint_dir}") - config = DiffusionModelConfig.from_pretrained( + config = DiffusionPipelineConfig.from_pretrained( checkpoint_dir, args=self.args, pipeline_config=resolved_pipeline_config, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 7be918d0ffec..4300a2943e80 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Pipeline registry for unified config flow. -Follows: VisualGenArgs → PipelineLoader → DiffusionModelConfig → AutoPipeline → BasePipeline +Follows: VisualGenArgs → PipelineLoader → DiffusionPipelineConfig → AutoPipeline → BasePipeline All pipelines (Wan, Flux, Flux2, LTX2, QwenImage) register via @register_pipeline decorator. @@ -35,7 +35,7 @@ from tensorrt_llm.logger import logger if TYPE_CHECKING: - from .config import DiffusionModelConfig + from .config import DiffusionPipelineConfig from .pipeline import BasePipeline @@ -123,11 +123,11 @@ class AutoPipeline: @staticmethod def from_config( - config: "DiffusionModelConfig", + config: "DiffusionPipelineConfig", checkpoint_dir: str, ) -> "BasePipeline": """ - Create pipeline instance from DiffusionModelConfig. + Create pipeline instance from DiffusionPipelineConfig. """ # Detect pipeline type from model_index.json or from model safetensors class_name = AutoPipeline._detect_from_checkpoint(checkpoint_dir) @@ -147,7 +147,7 @@ def from_config( logger.info(f"AutoPipeline: Creating {pipeline_class.__name__} from {checkpoint_dir}") - # Instantiate pipeline with DiffusionModelConfig + # Instantiate pipeline with DiffusionPipelineConfig return pipeline_class(config) @staticmethod diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8daefd5d7bb6..5b96985efd8f 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -481,7 +481,7 @@ class VisualGenArgs(StrictBaseModel): "Quantization config — accepts either a QuantConfig instance " "or a ModelOpt-format dict (e.g. ``{'quant_algo': 'FP8', " "'dynamic': True}``). Dict-form parsing happens lazily in " - "DiffusionModelConfig.from_pretrained." + "DiffusionPipelineConfig.from_pretrained." ), ) compilation_config: CompilationConfig = Field( diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index a5109b3cb01e..17ad0c2737ed 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -20,7 +20,7 @@ from test_common.llm_data import llm_models_root from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2_FORCE_ONE_STAGE_ENV from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, CacheDiTConfig, VisualGenArgs @@ -811,7 +811,7 @@ def test_resolve_variant_returns_two_stage_when_configured(self, monkeypatch): ) config = MagicMock() - config.pretrained_config._name_or_path = "" + config.primary_pretrained_config._name_or_path = "" config.extra_attrs = { "spatial_upsampler_path": "/fake/upsampler.safetensors", "distilled_lora_path": "/fake/lora.safetensors", @@ -828,7 +828,7 @@ def test_resolve_variant_honors_force_one_stage_env(self, monkeypatch): monkeypatch.setenv(LTX2_FORCE_ONE_STAGE_ENV, "1") config = MagicMock() - config.pretrained_config._name_or_path = "" + config.primary_pretrained_config._name_or_path = "" config.extra_attrs = { "spatial_upsampler_path": "/fake/upsampler.safetensors", "distilled_lora_path": "/fake/lora.safetensors", @@ -844,7 +844,7 @@ def test_resolve_variant_returns_base_without_two_stage_config(self): from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline config = MagicMock() - config.pretrained_config._name_or_path = "" + config.primary_pretrained_config._name_or_path = "" config.extra_attrs = {} result = LTX2Pipeline.resolve_variant(config) @@ -857,7 +857,7 @@ def test_resolve_variant_requires_both_paths(self): from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline config = MagicMock() - config.pretrained_config._name_or_path = "" + config.primary_pretrained_config._name_or_path = "" config.extra_attrs = {"spatial_upsampler_path": "/fake/upsampler.safetensors"} result = LTX2Pipeline.resolve_variant(config) @@ -881,7 +881,7 @@ def test_two_stage_auxiliary_paths_are_discovered_by_default(self, tmp_path, mon lora_path.touch() args = VisualGenArgs(model=str(checkpoint_path)) - config = DiffusionModelConfig.from_pretrained(str(checkpoint_path), args=args) + config = DiffusionPipelineConfig.from_pretrained(str(checkpoint_path), args=args) assert LTX2Pipeline.resolve_variant(config) is LTX2TwoStagesPipeline assert config.extra_attrs["spatial_upsampler_path"] == str(upsampler_path) @@ -898,7 +898,7 @@ def test_force_one_stage_env_skips_auto_discovery(self, tmp_path, monkeypatch): lora_path.touch() args = VisualGenArgs(model=str(checkpoint_path)) - config = DiffusionModelConfig.from_pretrained(str(checkpoint_path), args=args) + config = DiffusionPipelineConfig.from_pretrained(str(checkpoint_path), args=args) assert LTX2Pipeline.resolve_variant(config) is LTX2Pipeline assert "spatial_upsampler_path" not in config.extra_attrs @@ -919,7 +919,7 @@ def test_force_one_stage_env_prevents_promotion_with_explicit_auxiliary_paths( "distilled_lora_path": "/fake/lora.safetensors", }, ) - config = DiffusionModelConfig.from_pretrained(str(checkpoint_path), args=args) + config = DiffusionPipelineConfig.from_pretrained(str(checkpoint_path), args=args) assert config.extra_attrs["spatial_upsampler_path"] == "/fake/upsampler.safetensors" assert config.extra_attrs["distilled_lora_path"] == "/fake/lora.safetensors" @@ -941,7 +941,7 @@ def test_cache_dit_config_prevents_promotion_with_explicit_auxiliary_paths( "distilled_lora_path": "/fake/lora.safetensors", }, ) - config = DiffusionModelConfig.from_pretrained(str(checkpoint_path), args=args) + config = DiffusionPipelineConfig.from_pretrained(str(checkpoint_path), args=args) assert config.cache_backend == "cache_dit" assert config.extra_attrs["spatial_upsampler_path"] == "/fake/upsampler.safetensors" diff --git a/tests/unittest/_torch/visual_gen/test_model_loader.py b/tests/unittest/_torch/visual_gen/test_model_loader.py index e0e105b3bdf7..42bc14a52d9d 100644 --- a/tests/unittest/_torch/visual_gen/test_model_loader.py +++ b/tests/unittest/_torch/visual_gen/test_model_loader.py @@ -50,13 +50,13 @@ def test_meta_init_mode_creates_meta_tensors(checkpoint_exists): pytest.skip("Checkpoint not available") from tensorrt_llm._torch.models.modeling_utils import MetaInitMode - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig from tensorrt_llm._torch.visual_gen.models import AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs # Load config directly args = VisualGenArgs(model=CHECKPOINT_PATH) - config = DiffusionModelConfig.from_pretrained( + config = DiffusionPipelineConfig.from_pretrained( CHECKPOINT_PATH, args=args, ) @@ -174,15 +174,15 @@ def test_load_wan_pipeline_with_fp8_blockwise(checkpoint_exists): def test_visual_gen_args_to_quant_config(): """Test that VisualGenArgs accepts ModelOpt-format quant_config dicts. - The dict stays a dict on the public schema; DiffusionModelConfig + The dict stays a dict on the public schema; DiffusionPipelineConfig parses it (via load_diffusion_quant_config) when a pipeline loads. """ - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo from tensorrt_llm.visual_gen.args import VisualGenArgs - parse = DiffusionModelConfig.load_diffusion_quant_config + parse = DiffusionPipelineConfig.load_diffusion_quant_config # Default — no quantization. default_factory creates a QuantConfig # instance with quant_algo=None. @@ -268,7 +268,7 @@ def test_load_without_quant_config_no_fp8(checkpoint_exists): def test_visual_gen_args_from_dict(): """Test VisualGenArgs can be created from a dictionary.""" - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig from tensorrt_llm.quantization.mode import QuantAlgo from tensorrt_llm.visual_gen.args import VisualGenArgs @@ -283,7 +283,7 @@ def test_visual_gen_args_from_dict(): os.environ["WORLD_SIZE"] = "2" args = VisualGenArgs(**config_dict) assert args.model == "/path/to/model" - qc, _, dwq, _ = DiffusionModelConfig.load_diffusion_quant_config(args.quant_config) + qc, _, dwq, _ = DiffusionPipelineConfig.load_diffusion_quant_config(args.quant_config) assert qc.quant_algo == QuantAlgo.FP8 assert dwq is True assert args.parallel_config.ulysses_size == 2 diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py index d27ee61aa257..0a8ef1900479 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py @@ -10,7 +10,7 @@ # Importing the models package applies the Qwen-Image registration side effect. from tensorrt_llm._torch.visual_gen import models # noqa: F401 -from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenJointAttention from tensorrt_llm._torch.visual_gen.modules.attention import QKVMode from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader @@ -178,7 +178,7 @@ def test_qwen_pipeline_quant_config_parses_from_args( checkpoint_dir = _write_minimal_qwen_checkpoint(tmp_path) args = VisualGenArgs(model=str(checkpoint_dir), quant_config=quant_config) - config = DiffusionModelConfig.from_pretrained(str(checkpoint_dir), args=args) + config = DiffusionPipelineConfig.from_pretrained(str(checkpoint_dir), args=args) assert config.quant_config.quant_algo == quant_algo assert config.quant_config.group_size == group_size diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index a551c683696e..e5346b63b6bf 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -201,7 +201,7 @@ def test_nested_dict_auto_coerced(self): def test_quant_config_dict_passthrough(self): """ModelOpt-format dicts are accepted as-is — they parse in PipelineLoader.""" - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig raw = {"quant_algo": "FP8", "dynamic": True} args = VisualGenArgs(model="/tmp/model", quant_config=raw) @@ -210,8 +210,8 @@ def test_quant_config_dict_passthrough(self): assert isinstance(args.quant_config, dict) assert args.quant_config["quant_algo"] == "FP8" # The same dict is the source of truth for the derived flags; verify - # the parser DiffusionModelConfig.from_pretrained will run on it. - qc, _, dwq, daq = DiffusionModelConfig.load_diffusion_quant_config(args.quant_config) + # the pipeline-config parser will run on it. + qc, _, dwq, daq = DiffusionPipelineConfig.load_diffusion_quant_config(args.quant_config) assert qc.quant_algo is not None assert dwq is True diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py index 4c4a9d023a26..a274405b53e9 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py @@ -37,7 +37,11 @@ from diffusers import WanTransformer3DModel as HFWanTransformer3DModel from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, VisualGenArgs +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + DiffusionPipelineConfig, + VisualGenArgs, +) from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel from tensorrt_llm.models.modeling_utils import QuantConfig @@ -118,7 +122,9 @@ def _load_models(checkpoint_dir: str): ) args = VisualGenArgs(model=checkpoint_dir) - model_config = DiffusionModelConfig.from_pretrained(checkpoint_dir, args=args) + model_config = DiffusionPipelineConfig.from_pretrained(checkpoint_dir, args=args).model_configs[ + "transformer" + ] our_model = WanTransformer3DModel(model_config=model_config).to(DEVICE).eval() # Initialize our model with the exact same weights as the HF model. From 9bb57bfcb28ed05311f44f542a3de64ebc632715 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:21:46 +0800 Subject: [PATCH 2/3] [None][refactor] remove pipeline model_config alias Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../models/cosmos3/pipeline_cosmos3.py | 2 +- .../visual_gen/models/flux/pipeline_flux.py | 12 ++++--- .../visual_gen/models/flux/pipeline_flux2.py | 12 ++++--- .../visual_gen/models/ltx2/pipeline_ltx2.py | 18 +++++------ .../models/ltx2/pipeline_ltx2_two_stages.py | 6 ++-- .../models/qwen_image/pipeline_qwen_image.py | 8 ++--- .../visual_gen/models/wan/pipeline_wan.py | 18 ++++++----- .../visual_gen/models/wan/pipeline_wan_i2v.py | 18 ++++++----- tensorrt_llm/_torch/visual_gen/pipeline.py | 32 +++++++++---------- .../_torch/visual_gen/test_flux_pipeline.py | 10 +++--- .../_torch/visual_gen/test_ltx2_pipeline.py | 4 +-- .../_torch/visual_gen/test_model_loader.py | 4 +-- .../_torch/visual_gen/test_teacache.py | 11 ++++--- .../unittest/_torch/visual_gen/test_warmup.py | 4 +-- 14 files changed, 85 insertions(+), 74 deletions(-) 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 06b6f7fdad66..2977bcb64a61 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -68,7 +68,7 @@ def __init__(self, pipeline_config): def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") - self.transformer = Cosmos3VFMTransformer(self.model_configs["transformer"]) + self.transformer = Cosmos3VFMTransformer(self.pipeline_config.model_configs["transformer"]) def load_weights(self, weights: dict) -> None: if self.transformer is not None and hasattr(self.transformer, "load_weights"): diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index 883c8b3269e6..196c3c927e7a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -99,7 +99,7 @@ def _compute_flux_timestep_embedding( @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): @@ -121,7 +121,9 @@ def warmup_cache_key(self, height: int, width: int, **kwargs) -> tuple: def _init_transformer(self) -> None: """Initialize FLUX transformer with quantization support.""" logger.info("Creating FLUX transformer with quantization support...") - self.transformer = FluxTransformer2DModel(model_config=self.model_configs["transformer"]) + self.transformer = FluxTransformer2DModel( + model_config=self.pipeline_config.model_configs["transformer"] + ) def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): @@ -156,7 +158,7 @@ def load_standard_components( self.text_encoder = CLIPTextModel.from_pretrained( checkpoint_dir, subfolder=PipelineComponent.TEXT_ENCODER, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) # T5 tokenizer and text encoder (for sequence embeddings) @@ -171,7 +173,7 @@ def load_standard_components( self.text_encoder_2 = T5EncoderModel.from_pretrained( checkpoint_dir, subfolder=PipelineComponent.TEXT_ENCODER_2, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) # VAE @@ -203,7 +205,7 @@ def load_weights(self, weights: dict) -> None: self.transformer.load_weights(transformer_weights) logger.info("Transformer weights loaded successfully.") - self._target_dtype = self.model_config.torch_dtype + self._target_dtype = self.pipeline_config.torch_dtype if self.transformer is not None: self.transformer.eval() diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 81a40a50066a..82302aac8ce3 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -167,7 +167,7 @@ def _compute_flux2_timestep_embedding( @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): @@ -189,7 +189,9 @@ def warmup_cache_key(self, height: int, width: int, **kwargs) -> tuple: def _init_transformer(self) -> None: """Initialize FLUX.2 transformer with quantization support.""" logger.info("Creating FLUX.2 transformer with quantization support...") - self.transformer = Flux2Transformer2DModel(model_config=self.model_configs["transformer"]) + self.transformer = Flux2Transformer2DModel( + model_config=self.pipeline_config.model_configs["transformer"] + ) def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: with torch.no_grad(): @@ -255,13 +257,13 @@ def load_standard_components( # Mistral3 is a multimodal model (not pure CausalLM) self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained( text_encoder_path, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) else: # Qwen3 and other CausalLM text encoders self.text_encoder = AutoModelForCausalLM.from_pretrained( text_encoder_path, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) # VAE (FLUX.2-specific VAE with BatchNorm) @@ -294,7 +296,7 @@ def load_weights(self, weights: dict) -> None: self.transformer.load_weights(transformer_weights) logger.info("Transformer weights loaded successfully.") - self._target_dtype = self.model_config.torch_dtype + self._target_dtype = self.pipeline_config.torch_dtype if self.transformer is not None: self.transformer.eval() diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index ed25a96ed936..f48b3d05aea7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -644,7 +644,7 @@ def resolve_variant(cls, config): @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def default_warmup_resolutions(self): @@ -733,13 +733,13 @@ def _init_transformer(self) -> None: the reference ``LTXModelConfigurator.from_config()``. Missing keys fall back to the same defaults the reference uses. """ - attn_cfg = getattr(self.model_config, "attention", None) + attn_cfg = getattr(self.pipeline_config, "attention", None) if attn_cfg is not None and getattr(attn_cfg, "quant_attention_config", None) is not None: raise NotImplementedError( "Quantized attention is not yet supported for the LTX-2 pipeline." ) - model_config = self.model_configs["transformer"] + model_config = self.pipeline_config.model_configs["transformer"] cfg = model_config.pretrained_config rope_type = LTXRopeType(getattr(cfg, "rope_type", "interleaved")) @@ -807,11 +807,11 @@ def _setup_cuda_graphs(self): iterations (WARMUP_STEPS=2), so the captured graph contains the optimized compiled kernels. """ - if not self.model_config.cuda_graph.enable: + if not self.pipeline_config.cuda_graph.enable: return runner = _LTX2CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) - compile_note = " (with torch.compile)" if self.model_config.torch_compile.enable else "" + compile_note = " (with torch.compile)" if self.pipeline_config.torch_compile.enable else "" logger.info( f"CUDA graph runner: wrapping transformer.forward (Modality-aware){compile_note}" ) @@ -850,7 +850,7 @@ def load_standard_components( tokenizer files, and ``preprocessor_config.json``. """ skip_components = skip_components or [] - dtype = self.model_config.torch_dtype + dtype = self.pipeline_config.torch_dtype needs_text = ( PipelineComponent.TOKENIZER not in skip_components @@ -879,7 +879,7 @@ def load_standard_components( ).to(device) # --- Resolve native config ---------------------------------------- - native_config = self.model_config.extra_attrs.get("monolithic_safetensors_config") + native_config = self.pipeline_config.extra_attrs.get("monolithic_safetensors_config") sft_paths = _find_safetensors_files(checkpoint_dir) _prefetch_ltx2_safetensors_files(sft_paths) @@ -1021,7 +1021,7 @@ def post_load_weights(self) -> None: # self._setup_teacache(self.transformer, coefficients=LTX2_TEACACHE_COEFFICIENTS) # Cache-DiT - if self.transformer is not None and self.model_config.cache_backend == "cache_dit": + if self.transformer is not None and self.pipeline_config.cache_backend == "cache_dit": self._setup_cache_acceleration(self.transformer, coefficients=None) # Compression ratios from native scale factors @@ -1467,7 +1467,7 @@ def forward( # CFG parallel for multi-modal guidance: each GPU handles one # CFG pass (cond or uncond), results are all-gathered, then # STG/modality passes run on every GPU before the guidance formula. - vgm = self.model_config.visual_gen_mapping + vgm = self.pipeline_config.visual_gen_mapping cfg_size = vgm.cfg_size if vgm else 1 seq_parallel_size = vgm.seq_size if vgm is not None else 1 do_cfg_parallel_mm = use_multi_modal_guidance and cfg_size >= 2 and do_cfg diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 4b65cb8332e0..776ca9885a78 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -662,9 +662,9 @@ def load_standard_components( **kwargs, ) - dtype = self.model_config.torch_dtype - spatial_upsampler_path = self.model_config.extra_attrs.get("spatial_upsampler_path", "") - distilled_lora_path = self.model_config.extra_attrs.get("distilled_lora_path", "") + dtype = self.pipeline_config.torch_dtype + spatial_upsampler_path = self.pipeline_config.extra_attrs.get("spatial_upsampler_path", "") + distilled_lora_path = self.pipeline_config.extra_attrs.get("distilled_lora_path", "") # --- Spatial upsampler --- if spatial_upsampler_path: diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py index aed7e29554e9..1290453c57a0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py @@ -87,7 +87,7 @@ def __init__(self, pipeline_config): @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): @@ -122,7 +122,7 @@ def resolution_multiple_of(self) -> Tuple[int, int]: # ------------------------------------------------------------------ def _init_transformer(self) -> None: logger.info("Creating Qwen-Image transformer") - model_config = self.model_configs["transformer"] + model_config = self.pipeline_config.model_configs["transformer"] # ``pretrained_config`` is populated from # ``/transformer/config.json``. Read the fields we care # about with defaults matching the Qwen-Image 20B reference model. @@ -199,7 +199,7 @@ def load_standard_components( self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( checkpoint_dir, subfolder=PipelineComponent.TEXT_ENCODER, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) if PipelineComponent.VAE not in skip_components: @@ -232,7 +232,7 @@ def load_weights(self, weights: dict) -> None: # default. Cast only non-quantized tensors so FP8/NVFP4 weights # and FP32 scales keep the dtypes created by Linear.load_weights(). self.transformer.to_inference_dtype().eval() - self._target_dtype = self.model_config.torch_dtype + self._target_dtype = self.pipeline_config.torch_dtype # ------------------------------------------------------------------ # Prompt encoding (Qwen2.5-VL chat template). diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index c53dea7f0e13..7442f581f958 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -125,7 +125,7 @@ def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): t_emb = ce.time_embedder(t_freq) - teacache = self.model_config.teacache + teacache = self.pipeline_config.teacache if teacache is not None and teacache.use_ret_steps: return ce.time_proj(ce.act_fn(t_emb)).to(torch.float32) else: @@ -133,7 +133,7 @@ def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): @@ -177,13 +177,15 @@ def resolution_multiple_of(self): def _init_transformer(self) -> None: logger.info("Creating WAN transformer with quantization support...") - self.transformer = WanTransformer3DModel(model_config=self.model_configs["transformer"]) + self.transformer = WanTransformer3DModel( + model_config=self.pipeline_config.model_configs["transformer"] + ) # Wan2.2 A14B: create second transformer for two-stage denoising if self.is_wan22_14b: logger.info("Creating second transformer for Wan2.2 A14B two-stage denoising...") self.transformer_2 = WanTransformer3DModel( - model_config=self.model_configs["transformer_2"] + model_config=self.pipeline_config.model_configs["transformer_2"] ) def load_standard_components( @@ -225,7 +227,7 @@ def load_standard_components( self.text_encoder = UMT5EncoderModel.from_pretrained( checkpoint_dir, subfolder=PipelineComponent.TEXT_ENCODER, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) if PipelineComponent.VAE not in skip_components: @@ -289,7 +291,7 @@ def load_weights(self, weights: dict) -> None: logger.info("Transformer_2 weights loaded successfully.") # Cache the target dtype from model config (default: bfloat16) - self._target_dtype = self.model_config.torch_dtype + self._target_dtype = self.pipeline_config.torch_dtype # Set model to eval mode if self.transformer is not None: @@ -301,7 +303,7 @@ def post_load_weights(self) -> None: super().post_load_weights() # Calls transformer.post_load_weights() for FP8 scale transformations if self.transformer is not None: # TeaCache extractor only when using TeaCache (not Cache-DiT). - if self.model_config.cache_backend == "teacache": + if self.pipeline_config.cache_backend == "teacache": register_extractor_from_config( ExtractorConfig( model_class_name="WanTransformer3DModel", @@ -316,7 +318,7 @@ def post_load_weights(self) -> None: ) self.transformer_cache_backend = self.cache_accelerator else: - if self.model_config.cache_backend == "cache_dit": + if self.pipeline_config.cache_backend == "cache_dit": self._setup_cache_acceleration(self.transformer, coefficients=None) # TeaCache is not supported for Wan 2.2 unless using Cache-DiT. self.transformer_cache_backend = self.cache_accelerator diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index dd70e3005612..2ff09e153a03 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -121,7 +121,7 @@ def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): t_emb = ce.time_embedder(t_freq) - teacache = self.model_config.teacache + teacache = self.pipeline_config.teacache if teacache is not None and teacache.use_ret_steps: # ret_steps mode: use timestep_proj — what the ret_steps coefficients were calibrated for return ce.time_proj(ce.act_fn(t_emb)).to(torch.float32) @@ -130,7 +130,7 @@ def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): @@ -168,13 +168,15 @@ def resolution_multiple_of(self): def _init_transformer(self) -> None: logger.info("Creating WAN I2V transformer with quantization support...") - self.transformer = WanTransformer3DModel(model_config=self.model_configs["transformer"]) + self.transformer = WanTransformer3DModel( + model_config=self.pipeline_config.model_configs["transformer"] + ) # Wan2.2: Optionally create second transformer for two-stage denoising if self.boundary_ratio is not None: logger.info("Creating second transformer for Wan2.2 I2V two-stage denoising...") self.transformer_2 = WanTransformer3DModel( - model_config=self.model_configs["transformer_2"] + model_config=self.pipeline_config.model_configs["transformer_2"] ) def load_standard_components( @@ -221,7 +223,7 @@ def load_standard_components( self.text_encoder = UMT5EncoderModel.from_pretrained( checkpoint_dir, subfolder=PipelineComponent.TEXT_ENCODER, - torch_dtype=self.model_config.torch_dtype, + torch_dtype=self.pipeline_config.torch_dtype, ).to(device) if PipelineComponent.VAE not in skip_components: @@ -310,7 +312,7 @@ def load_weights(self, weights: dict) -> None: logger.info("Transformer_2 weights loaded successfully.") # Cache the target dtype from model config (default: bfloat16) - self._target_dtype = self.model_config.torch_dtype + self._target_dtype = self.pipeline_config.torch_dtype # Set model to eval mode if self.transformer is not None: @@ -323,7 +325,7 @@ def load_weights(self, weights: dict) -> None: def post_load_weights(self) -> None: super().post_load_weights() # Calls transformer.post_load_weights() for FP8 scale transformations if self.transformer is not None: - if self.model_config.cache_backend == "teacache": + if self.pipeline_config.cache_backend == "teacache": register_extractor_from_config( ExtractorConfig( model_class_name="WanTransformer3DModel", @@ -338,7 +340,7 @@ def post_load_weights(self) -> None: ) self.transformer_cache_backend = self.cache_accelerator else: - if self.model_config.cache_backend == "cache_dit": + if self.pipeline_config.cache_backend == "cache_dit": self._setup_cache_acceleration(self.transformer, coefficients=None) self.transformer_cache_backend = self.cache_accelerator diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 10e87d3fe43d..7a0c629bf5c9 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -120,8 +120,6 @@ def resolve_variant(cls, config: "DiffusionPipelineConfig") -> Type["BasePipelin def __init__(self, pipeline_config: "DiffusionPipelineConfig"): super().__init__() self.pipeline_config = pipeline_config - self.model_config = pipeline_config - self.model_configs = pipeline_config.model_configs self.config = pipeline_config.primary_pretrained_config self.mapping: Mapping = getattr(pipeline_config, "mapping", None) or Mapping() self._cuda_graph_runners: Dict[str, CUDAGraphRunner] = {} @@ -173,10 +171,10 @@ def _cuda_profiler_stop(self): def _setup_cuda_graphs(self): """Wrap all transformer components with CUDA graph capture/replay.""" - if not self.model_config.cuda_graph.enable: + if not self.pipeline_config.cuda_graph.enable: return - if self.model_config.torch_compile.enable: + if self.pipeline_config.torch_compile.enable: logger.warning( "CUDA graphs with torch.compile not yet supported. Using torch.compile only." ) @@ -296,7 +294,7 @@ def resolve_warmup_plan(self) -> Tuple[List[Tuple[int, int, int]], int]: Returns: (shapes, steps) tuple where shapes = list of (h, w, f) """ - warmup_cfg = self.model_config.compilation + warmup_cfg = self.pipeline_config.compilation if warmup_cfg.resolutions is not None or warmup_cfg.num_frames is not None: resolutions = ( @@ -399,11 +397,13 @@ def post_load_weights(self) -> None: self.transformer.post_load_weights() def _apply_teacache_coefficients(self, coefficients: Optional[Dict]) -> None: - """Pick TeaCache coefficients from checkpoint path; updates model_config.teacache in place.""" + """Pick TeaCache coefficients from checkpoint path; updates pipeline config in place.""" if not coefficients: return - teacache_cfg = self.model_config.teacache - checkpoint_path = getattr(self.model_config.primary_pretrained_config, "_name_or_path", "") + teacache_cfg = self.pipeline_config.teacache + checkpoint_path = getattr( + self.pipeline_config.primary_pretrained_config, "_name_or_path", "" + ) matched = False for model_size, coeff_data in coefficients.items(): if model_size.lower() in checkpoint_path.lower(): @@ -444,7 +444,7 @@ def _setup_cache_acceleration( self.cache_accelerator.unwrap() self.cache_accelerator = None - cfg = self.model_config + cfg = self.pipeline_config if cfg.cache_backend == "cache_dit": acc = CacheDiTAccelerator(self, cfg.cache_dit) @@ -475,8 +475,8 @@ def setup_parallel_vae(self): parallel-VAE decode ownership applies. The actual ``ParallelVAEFactory`` wrap is a local side effect that only runs on ranks in ``vae_ranks``. """ - parallel_cfg = self.model_config.parallel - vgm = self.model_config.visual_gen_mapping + parallel_cfg = self.pipeline_config.parallel + vgm = self.pipeline_config.visual_gen_mapping # Global preconditions — evaluate identically on every rank. self._parallel_vae_enabled = ( @@ -527,7 +527,7 @@ def torch_compile(self) -> None: For non-transformer components, compiles the entire module. """ - tc_config = self.model_config.torch_compile + tc_config = self.pipeline_config.torch_compile # Using default as max-autotune mode takes more initialization time and # does not improve performance a lot. @@ -662,7 +662,7 @@ def decode_latents( Non-decoding ranks return ``None`` (or a tuple of ``None``). """ if self._parallel_vae_enabled: - vgm = self.model_config.visual_gen_mapping + vgm = self.pipeline_config.visual_gen_mapping decode_ranks = set(vgm.vae_ranks) else: decode_ranks = {0} @@ -702,7 +702,7 @@ def _setup_cfg_config( Returns: Dict with CFG configuration including split tensors """ - vgm = self.model_config.visual_gen_mapping + vgm = self.pipeline_config.visual_gen_mapping cfg_size = vgm.cfg_size if vgm else 1 ulysses_size = vgm.ulysses_size if vgm else 1 attn2d_row_size = vgm.attn2d_row_size if vgm else 1 @@ -773,7 +773,7 @@ def _denoise_step_cfg_parallel( local_extras, ): """Execute single denoising step with CFG parallel.""" - vgm = self.model_config.visual_gen_mapping + vgm = self.pipeline_config.visual_gen_mapping cfg_pg = vgm.cfg_group if vgm else None cfg_size = vgm.cfg_size if vgm else 1 @@ -1105,7 +1105,7 @@ def denoise( if getattr(self, "cache_accelerator", None) and self.cache_accelerator.is_enabled(): stats = self.cache_accelerator.get_stats() if stats: - if self.model_config.cache_backend == "cache_dit": + if self.pipeline_config.cache_backend == "cache_dit": logger.info("Cache-DiT stats: %s", stats) elif "hit_rate" in stats: logger.info( diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py index c51c338af2fe..0455e25faff7 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py @@ -176,7 +176,7 @@ def test_load_flux1_pipeline_basic(self, flux1_checkpoint_exists): assert pipeline is not None assert hasattr(pipeline, "transformer") assert pipeline.transformer is not None - assert pipeline.model_config.attention.backend == "VANILLA" + assert pipeline.pipeline_config.attention.backend == "VANILLA" del pipeline gc.collect() @@ -210,7 +210,7 @@ def test_load_flux1_with_attention_backend(self, flux1_checkpoint_exists, backen pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) - assert pipeline.model_config.attention.backend == backend + assert pipeline.pipeline_config.attention.backend == backend del pipeline gc.collect() @@ -236,7 +236,7 @@ def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo: pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) - assert pipeline.model_config.quant_config.quant_algo is not None + assert pipeline.pipeline_config.quant_config.quant_algo is not None # Count quantized Linear layers and verify FP8 weights quant_count = 0 @@ -277,7 +277,7 @@ def test_load_flux2_with_quantization(self, flux2_checkpoint_exists, quant_algo: pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) - assert pipeline.model_config.quant_config.quant_algo is not None + assert pipeline.pipeline_config.quant_config.quant_algo is not None quant_count = 0 found_fp8 = False @@ -1157,7 +1157,7 @@ def _run_all_optimizations_worker( transformer = pipeline.transformer.eval() # Verify all optimizations are enabled - assert pipeline.model_config.visual_gen_mapping.ulysses_size == world_size, ( + assert pipeline.pipeline_config.visual_gen_mapping.ulysses_size == world_size, ( "Ulysses parallel not enabled" ) assert transformer.model_config.quant_config.quant_algo == QuantAlgo.FP8, "FP8 not enabled" diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index 17ad0c2737ed..8fc668929427 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -200,7 +200,7 @@ def test_load_with_quantization(self, ltx2_bf16_checkpoint_exists, quant_algo: s pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) - assert pipeline.model_config.quant_config.quant_algo is not None + assert pipeline.pipeline_config.quant_config.quant_algo is not None quant_count = 0 found_fp8 = False @@ -1205,7 +1205,7 @@ def test_two_stage_with_quantization(self, ltx2_two_stage_assets_exist, quant_al pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) try: assert isinstance(pipeline, LTX2TwoStagesPipeline) - assert pipeline.model_config.quant_config.quant_algo is not None + assert pipeline.pipeline_config.quant_config.quant_algo is not None quant_count = sum( 1 diff --git a/tests/unittest/_torch/visual_gen/test_model_loader.py b/tests/unittest/_torch/visual_gen/test_model_loader.py index 42bc14a52d9d..e4bb4de9ee84 100644 --- a/tests/unittest/_torch/visual_gen/test_model_loader.py +++ b/tests/unittest/_torch/visual_gen/test_model_loader.py @@ -124,7 +124,7 @@ def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists): pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) # Verify model config has dynamic_weight_quant enabled - assert pipeline.model_config.dynamic_weight_quant is True, ( + assert pipeline.pipeline_config.dynamic_weight_quant is True, ( "dynamic_weight_quant should be True when linear.type specifies FP8" ) @@ -252,7 +252,7 @@ def test_load_without_quant_config_no_fp8(checkpoint_exists): pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) # Verify dynamic_weight_quant is False - assert pipeline.model_config.dynamic_weight_quant is False, ( + assert pipeline.pipeline_config.dynamic_weight_quant is False, ( "dynamic_weight_quant should be False when no quant_config" ) diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index c65326cf8c47..a0080a6d2008 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -20,7 +20,7 @@ import pytest from tensorrt_llm._torch.visual_gen.cache.teacache import TeaCacheBackend -from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm.visual_gen.args import TeaCacheConfig @@ -31,8 +31,11 @@ class TestSetupTeacache: def _make_pipeline_mock(self, checkpoint_name, use_ret_steps=False): pipeline = MagicMock() pipeline.cache_accelerator = None - pipeline.model_config = DiffusionModelConfig( + model_config = DiffusionModelConfig( pretrained_config=SimpleNamespace(_name_or_path=f"/path/to/{checkpoint_name}/snapshot"), + ) + pipeline.pipeline_config = DiffusionPipelineConfig( + model_configs={"transformer": model_config}, cache=TeaCacheConfig( teacache_thresh=0.3, use_ret_steps=use_ret_steps, @@ -51,7 +54,7 @@ def test_matching_variant_selects_coefficients(self): with patch.object(TeaCacheBackend, "enable"): BasePipeline._setup_cache_acceleration(pipeline, MagicMock(), coefficients) - assert pipeline.model_config.teacache.coefficients == [1.0, 2.0, 3.0] + assert pipeline.pipeline_config.teacache.coefficients == [1.0, 2.0, 3.0] def test_no_match_raises_valueerror(self): """Raises ValueError (fail-early) when no variant matches checkpoint.""" @@ -66,7 +69,7 @@ def test_no_match_raises_valueerror(self): def test_disabled_teacache_is_noop(self): """No-op when cache is None (TeaCache not selected).""" pipeline = self._make_pipeline_mock("FLUX.1-dev") - pipeline.model_config = pipeline.model_config.model_copy(update={"cache": None}) + pipeline.pipeline_config = pipeline.pipeline_config.model_copy(update={"cache": None}) BasePipeline._setup_cache_acceleration(pipeline, MagicMock(), {"dev": [1.0]}) assert pipeline.cache_accelerator is None diff --git a/tests/unittest/_torch/visual_gen/test_warmup.py b/tests/unittest/_torch/visual_gen/test_warmup.py index f7b6672df041..c882da486f70 100644 --- a/tests/unittest/_torch/visual_gen/test_warmup.py +++ b/tests/unittest/_torch/visual_gen/test_warmup.py @@ -102,8 +102,8 @@ class _BaseStubPipeline(BasePipeline): def __init__(self, warmup_cfg): self._warmed_up_shapes = set() - self.model_config = MagicMock() - self.model_config.compilation = warmup_cfg or CompilationConfig() + self.pipeline_config = MagicMock() + self.pipeline_config.compilation = warmup_cfg or CompilationConfig() def forward(self, *args, **kwargs): pass From fde2e27a87f34f2b9943242e52d1294995f1b76e Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:27:45 +0800 Subject: [PATCH 3/3] [None][test] fix visual gen args review lint Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tests/unittest/_torch/visual_gen/test_visual_gen_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index e5346b63b6bf..540cc0c7f4c8 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -211,7 +211,7 @@ def test_quant_config_dict_passthrough(self): assert args.quant_config["quant_algo"] == "FP8" # The same dict is the source of truth for the derived flags; verify # the pipeline-config parser will run on it. - qc, _, dwq, daq = DiffusionPipelineConfig.load_diffusion_quant_config(args.quant_config) + qc, _, dwq, _ = DiffusionPipelineConfig.load_diffusion_quant_config(args.quant_config) assert qc.quant_algo is not None assert dwq is True