fix(native_transformer): forward load kwargs to from_config

load_transformer threads **kwargs into its from_pretrained and from_single_file branches. The native dispatch branch bypasses both loaders and builds via cls.from_config, so those kwargs were dropped there. Thread them through load, build_component, and build_component_quantized into from_config so callers passing extra args alongside native_spec are honored instead of silently dropped. Siblings do not receive them.

pipelines/generic.py

@@ -76,6 +76,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
                 modules_dtype_dict=modules_dtype_dict,
                 quant_args=quant_args,
                 quant_type=quant_type,
+                **kwargs,
             )
         elif local_file is not None and local_file.lower().endswith('.safetensors'):
             log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}')

pipelines/native_transformer.py

@@ -179,6 +179,7 @@ def load(
     modules_dtype_dict: dict | None = None,
     quant_args: dict | None = None,
     quant_type: str | None = None,
+    **kwargs,
 ) -> tuple[object, dict[str, object]]:
     """Load the transformer (and any bundled siblings) from ``local_file``.
 
@@ -193,7 +194,11 @@ def load(
     :func:`pipelines.generic.load_transformer` so the dispatch from there can
     plumb the caller's intent through unchanged. ``quant_args`` and
     ``quant_type`` are precomputed by the caller; when ``None`` they are
-    derived here via ``model_quant.get_dit_args``.
+    derived here via ``model_quant.get_dit_args``. Extra ``**kwargs`` are
+    forwarded to the transformer's ``cls.from_config``, the native path's
+    construction step (it bypasses ``from_pretrained``/``from_single_file``,
+    where ``load_transformer`` otherwise routes them). Siblings do not
+    receive them.
 
     Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by
     sibling name and is empty for non-sibling specs, or for sibling specs
@@ -247,6 +252,7 @@ def load(
         dtype=effective_dtype,
         modules_to_not_convert=modules_to_not_convert,
         modules_dtype_dict=modules_dtype_dict,
+        **kwargs,
     )
     del transformer_sd
     devices.torch_gc()
@@ -386,6 +392,7 @@ def build_component_quantized(
     quant_args: dict,
     dtype,
     acceptable_missing: tuple[str, ...],
+    **kwargs,
 ) -> object:
     """Build a component with per-tensor SDNQ quantization during load.
 
@@ -424,7 +431,7 @@ def build_component_quantized(
     quantizer.torch_dtype = target_dtype
 
     with init_empty_weights(include_buffers=False):
-        component = cls.from_config(config)
+        component = cls.from_config(config, **kwargs)
 
     quantizer._process_model_before_weight_loading(component, device_map=None)  # pylint: disable=protected-access
 
@@ -483,6 +490,7 @@ def build_component(
     dtype=None,
     modules_to_not_convert: list | None = None,
     modules_dtype_dict: dict | None = None,
+    **kwargs,
 ) -> object:
     """Convert (if needed), instantiate, load weights, dtype-cast, quantize,
     and offload-place a single component. Raises on any hard failure.
@@ -496,7 +504,8 @@ def build_component(
     ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global
     default is used. ``modules_to_not_convert`` and ``modules_dtype_dict``
     are forwarded to :func:`apply_quant` for the post-mode path; pre-mode
-    receives them via the SDNQConfig in ``quant_args``.
+    receives them via the SDNQConfig in ``quant_args``. Extra ``**kwargs``
+    reach ``cls.from_config`` for both construction paths.
     """
     try:
         if converter is not None:
@@ -514,13 +523,14 @@ def build_component(
                 quant_args=quant_args,
                 dtype=dtype,
                 acceptable_missing=acceptable_missing,
+                **kwargs,
             )
             del sd
             devices.torch_gc()
             return component
 
         log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}')
-        component = cls.from_config(config)
+        component = cls.from_config(config, **kwargs)
         missing, unexpected = component.load_state_dict(sd, strict=False)
         validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
         del sd

test/test-native-transformer.py

@@ -422,6 +422,19 @@ class MockMiniTransformer(torch.nn.Module):
         self.rope.register_buffer('freqs', torch.zeros(dim))
 
 
+class MockKwargsTransformer(MockMiniTransformer):
+    """Records the kwargs from_config received. Mirrors diffusers from_config,
+    which accepts **kwargs (config overrides); lets a test assert the native
+    load path forwards caller kwargs to construction instead of dropping them."""
+
+    last_kwargs: dict = {}
+
+    @classmethod
+    def from_config(cls, config: dict, **kwargs) -> 'MockKwargsTransformer':
+        cls.last_kwargs = dict(kwargs)
+        return cls(dim=config['dim'])
+
+
 def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str:
     os.close(fd)
     safetensors.torch.save_file(state_dict_keys, path)
@@ -484,6 +497,55 @@ def test_load_end_to_end_with_bfl_prefix_no_converter():
             os.unlink(path)
 
 
+def test_load_forwards_kwargs_to_from_config():
+    """Caller **kwargs reach cls.from_config (the native path's only
+    construction step, since it bypasses from_pretrained/from_single_file).
+    Guards against silently dropping args a caller passes alongside
+    native_spec."""
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim = 8
+        raw = {
+            'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
+            'model.diffusion_model.in_proj.bias': torch.zeros(dim),
+            'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
+            'model.diffusion_model.out_proj.bias': torch.zeros(dim),
+        }
+        write_fixture(raw, fd, path)
+
+        orig_fetch = nt.fetch_component_config
+        nt.fetch_component_config = lambda repo, sub: {'dim': dim}
+        from modules import model_quant
+        orig_get_dit = model_quant.get_dit_args
+        orig_get_qtype = model_quant.get_quant_type
+        orig_do_post = model_quant.do_post_load_quant
+        model_quant.get_dit_args = lambda *a, **k: ({}, {})
+        model_quant.get_quant_type = lambda *a, **k: None
+        model_quant.do_post_load_quant = lambda *a, **k: None
+
+        MockKwargsTransformer.last_kwargs = {}
+        try:
+            spec = nt.TransformerSpec(cls=MockKwargsTransformer)
+            transformer, _ = nt.load(
+                local_file=path,
+                repo_id='fake/repo',
+                spec=spec,
+                diffusers_cfg={},
+                low_cpu_mem_usage=True,
+            )
+        finally:
+            nt.fetch_component_config = orig_fetch
+            model_quant.get_dit_args = orig_get_dit
+            model_quant.get_quant_type = orig_get_qtype
+            model_quant.do_post_load_quant = orig_do_post
+
+        assert MockKwargsTransformer.last_kwargs == {'low_cpu_mem_usage': True}
+        assert isinstance(transformer, MockKwargsTransformer)
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
 def test_load_end_to_end_with_sibling_partition():
     """Bundled-sibling case: file carries both transformer and sibling weights,
     sibling_classes supplies the runtime sibling class, partition routes each
@@ -711,6 +773,7 @@ def run_all():
     cat = category('load')
     for fn in [
         test_load_end_to_end_with_bfl_prefix_no_converter,
+        test_load_forwards_kwargs_to_from_config,
         test_load_end_to_end_with_sibling_partition,
         test_load_raises_on_missing_sibling_class,
         test_load_rejects_non_safetensors,