[None][refactor] split VisualGen pipeline and model configs

tensorrt_llm/_torch/visual_gen/__init__.py

@@ -12,14 +12,15 @@
 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
 from .pipeline_registry import PipelineComponent
 
 __all__ = [
     "DiffusionModelConfig",
+    "DiffusionPipelineConfig",
     "PipelineComponent",
     "WeightLoader",
     "PipelineLoader",

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

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.pipeline_config.model_configs["transformer"])
 
     def load_weights(self, weights: dict) -> None:
         if self.transformer is not None and hasattr(self.transformer, "load_weights"):

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

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(
@@ -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_config)
+        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()