fix(krea2): zero-fill dormant last.up/last.down residual branch

Base ships last.up all-zeros so the branch is a no-op; pre-branch
finetunes omit both keys. A zero_init_missing spec field zero-fills
them on load instead of falling back to the base transformer.

pipelines/krea2/__init__.py

@@ -5,4 +5,11 @@ from pipelines.native_transformer import TransformerSpec
 # Checkpoint keys are bare (`first.`, `blocks.N.`, `txtfusion.`, ...) and the transformer's
 # module tree mirrors them exactly, so no state-dict conversion is needed. The model has no
 # rope/pos buffers, so the default acceptable-missing set is sufficient.
-KREA2_SPEC = TransformerSpec(cls=Krea2Transformer2DModel, converter=None)
+#
+# `last.up`/`last.down` are a dormant zero-init residual branch (base ships last.up.weight
+# all-zeros). Finetunes predating it omit both keys; zero-fill so the branch stays a no-op.
+KREA2_SPEC = TransformerSpec(
+    cls=Krea2Transformer2DModel,
+    converter=None,
+    zero_init_missing=("last.down.", "last.up."),
+)

pipelines/native_transformer.py

@@ -100,6 +100,12 @@ class TransformerSpec:
     prefixes, no converter, no siblings). Arches with bundled-sibling
     components (Anima) or unusual key conventions (custom converters,
     Cosmos-style structural markers) override the relevant fields.
+
+    ``zero_init_missing`` lists key prefixes for a zero-initialized residual
+    branch the class defines but some checkpoints omit. Unlike
+    ``acceptable_missing`` (buffer-only keys left at their init), these are
+    zero-filled on load so the branch stays a no-op, matching a base model
+    that ships the branch dormant (output projection all-zeros).
     """
 
     cls: type
@@ -108,6 +114,7 @@ class TransformerSpec:
     converter: Callable[[dict], dict] | None = None
     siblings: dict[str, SiblingSpec] = field(default_factory=dict)
     acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING
+    zero_init_missing: tuple[str, ...] = ()
     forbidden_markers: tuple[tuple[str, str], ...] = ()
 
 
@@ -260,6 +267,7 @@ def load(
         cls=spec.cls,
         converter=spec.converter,
         acceptable_missing=spec.acceptable_missing,
+        zero_init_missing=spec.zero_init_missing,
         quant_args=quant_args,
         quant_type=quant_type,
         dtype=effective_dtype,
@@ -401,6 +409,7 @@ def build_component_quantized(
     quant_args: dict,
     dtype,
     acceptable_missing: tuple[str, ...],
+    zero_init_missing: tuple[str, ...] = (),
     **kwargs,
 ) -> object:
     """Build a component with per-tensor SDNQ quantization during load.
@@ -480,6 +489,9 @@ def build_component_quantized(
             pbar.update(task, advance=1)
 
     missing = sorted(expected_keys - loaded_keys)
+    missing = materialize_zero_init(
+        component, missing, zero_init_missing, device=target_device, dtype=target_dtype
+    )
     validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
 
     component = quantizer._process_model_after_weight_loading(component)  # pylint: disable=protected-access
@@ -494,6 +506,7 @@ def build_component(
     cls: type,
     converter: Callable[[dict], dict] | None,
     acceptable_missing: tuple[str, ...],
+    zero_init_missing: tuple[str, ...] = (),
     quant_args: dict,
     quant_type: str | None,
     dtype=None,
@@ -539,6 +552,7 @@ def build_component(
                 quant_args=quant_args,
                 dtype=dtype,
                 acceptable_missing=acceptable_missing,
+                zero_init_missing=zero_init_missing,
                 **kwargs,
             )
             del sd
@@ -548,6 +562,7 @@ def build_component(
         log.debug(f'Load model: transformer=native {component_name} loading keys={len(sd)} cls={cls.__name__}')
         component = cls.from_config(config, **kwargs)
         missing, unexpected = component.load_state_dict(sd, strict=False)
+        missing = materialize_zero_init(component, missing, zero_init_missing)
         validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
         del sd
         devices.torch_gc()
@@ -613,6 +628,57 @@ def validate_state_dict_load(
         )
 
 
+def materialize_zero_init(
+    component: object,
+    missing: list[str],
+    zero_init_missing: tuple[str, ...],
+    *,
+    device=None,
+    dtype=None,
+) -> list[str]:
+    """Zero-fill missing weights that belong to a zero-initialized residual
+    branch the class defines but the override file predates.
+
+    A base model may ship a residual branch dormant: its output projection is
+    all-zeros, so ``out + up(down(x))`` reduces to ``out`` and the branch is a
+    no-op. A checkpoint made before the branch existed omits those keys. Left
+    unmaterialized they keep random ``from_config`` init (non-quant path) or
+    stay on the meta device (per-tensor quant path); either injects garbage
+    into the output. Zero-filling reproduces the base's dormant behavior.
+
+    Returns ``missing`` with the handled keys removed, so
+    :func:`validate_state_dict_load` does not treat them as a hard mismatch.
+    ``device``/``dtype`` override the target placement (the quant path passes
+    the real device since its params are on meta); when omitted each key's own
+    parameter device/dtype is used. Keys matching a prefix but absent from the
+    component's parameters are left in ``missing`` untouched.
+    """
+    if not zero_init_missing:
+        return missing
+    from accelerate.utils import set_module_tensor_to_device
+    params = dict(component.named_parameters())
+    handled: list[str] = []
+    remaining: list[str] = []
+    for key in missing:
+        param = params.get(key)
+        if param is not None and any(key.startswith(p) for p in zero_init_missing):
+            tgt_device = device if device is not None else param.device
+            tgt_dtype = dtype if dtype is not None else param.dtype
+            set_module_tensor_to_device(
+                component, key, tgt_device,
+                value=torch.zeros(param.shape, dtype=tgt_dtype), dtype=tgt_dtype,
+            )
+            handled.append(key)
+        else:
+            remaining.append(key)
+    if handled:
+        log.debug(
+            f"Load model: transformer=native zero-init {len(handled)} dormant "
+            f"residual key(s): {', '.join(handled)}"
+        )
+    return remaining
+
+
 def apply_quant(
     transformer: object,
     quant_type: str | None,

test/test-krea2-transformer.py

@@ -1,11 +1,15 @@
 #!/usr/bin/env python
-"""Offline parity test: Krea2Transformer2DModel vs the reference SingleStreamDiT.
+"""Offline tests for the Krea2 transformer port and its native loader handling.
 
-Builds both models from one tiny config, copies the reference state dict into the diffusers
-port, runs identical inputs, and asserts the forward outputs match. No server, no checkpoint.
+- Parity: Krea2Transformer2DModel vs the reference SingleStreamDiT. Builds both from one tiny
+  config, copies the reference state dict into the diffusers port, runs identical inputs, and
+  asserts the forward outputs match. Needs the reference (mmdit.py) at $KREA2_REF_DIR
+  (default /home/ohiom/database/watering-hole).
+- Zero-init regression: a checkpoint that omits the dormant last.up/last.down residual branch,
+  after materialize_zero_init, produces the same output as the base whose up is zeroed. Port
+  only, so it runs without the reference.
 
-The reference checkpoint repo (mmdit.py) is expected at $KREA2_REF_DIR
-(default /home/ohiom/database/watering-hole).
+No server, no checkpoint.
 """
 
 import importlib.util
@@ -17,6 +21,12 @@ import torch
 
 REF_DIR = os.environ.get("KREA2_REF_DIR", "/home/ohiom/database/watering-hole")
 
+CFG = dict(
+    features=128, tdim=32, txtdim=64, heads=4, kvheads=2, multiplier=4,
+    layers=2, patch=2, channels=4, bias=False, theta=1e3,
+    txtlayers=3, txtheads=2, txtkvheads=2,
+)
+
 
 def load_reference():
     sys.path.insert(0, REF_DIR)
@@ -35,34 +45,29 @@ def load_port():
     return mod
 
 
-def main():
-    mmdit = load_reference()
-    port = load_port()
-
-    cfg = dict(
-        features=128, tdim=32, txtdim=64, heads=4, kvheads=2, multiplier=4,
-        layers=2, patch=2, channels=4, bias=False, theta=1e3,
-        txtlayers=3, txtheads=2, txtkvheads=2,
-    )
-
-    torch.manual_seed(0)
-    ref = mmdit.SingleStreamDiT(mmdit.SingleMMDiTConfig(**cfg)).float().eval()
-    mine = port.Krea2Transformer2DModel(**cfg).float().eval()
-    missing, unexpected = mine.load_state_dict(ref.state_dict(), strict=False)
-    assert not missing, f"missing keys when loading reference weights: {missing}"
-    assert not unexpected, f"unexpected keys when loading reference weights: {unexpected}"
-
+def make_inputs():
     batch, txtlen, imglen = 2, 5, 9
-    cdim = cfg["channels"] * cfg["patch"] ** 2
+    cdim = CFG["channels"] * CFG["patch"] ** 2
     seq = txtlen + imglen
     gen = torch.Generator().manual_seed(1)
     img = torch.randn(batch, imglen, cdim, generator=gen)
-    context = torch.randn(batch, txtlen, cfg["txtlayers"], cfg["txtdim"], generator=gen)
+    context = torch.randn(batch, txtlen, CFG["txtlayers"], CFG["txtdim"], generator=gen)
     timestep = torch.rand(batch, generator=gen)
     pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float()
     mask = torch.ones(batch, seq, dtype=torch.bool)
     mask[0, -2:] = False  # exercise the key-padding path
+    return img, context, timestep, pos, mask
 
+
+def run_parity(mmdit, port):
+    torch.manual_seed(0)
+    ref = mmdit.SingleStreamDiT(mmdit.SingleMMDiTConfig(**CFG)).float().eval()
+    mine = port.Krea2Transformer2DModel(**CFG).float().eval()
+    missing, unexpected = mine.load_state_dict(ref.state_dict(), strict=False)
+    assert not missing, f"missing keys when loading reference weights: {missing}"
+    assert not unexpected, f"unexpected keys when loading reference weights: {unexpected}"
+
+    img, context, timestep, pos, mask = make_inputs()
     with torch.no_grad():
         out_ref = ref(img, context, timestep, pos, mask)
         out_mine = mine(
@@ -84,5 +89,83 @@ def main():
     print("PARITY OK")
 
 
+def load_materialize_zero_init():
+    """Import the real loader helper. native_transformer pulls in modules.shared, which needs
+    cmd_args parsed first, so bootstrap it the same way the native-transformer suite does."""
+    repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    if repo not in sys.path:
+        sys.path.insert(0, repo)
+    os.environ.setdefault("SD_INSTALL_QUIET", "1")
+    import modules.cmd_args
+    import installer
+    orig_argv = sys.argv
+    sys.argv = [sys.argv[0]]
+    try:
+        modules.cmd_args.parse_args()
+    finally:
+        sys.argv = orig_argv
+    installer.add_args(modules.cmd_args.parser)
+    modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
+    from pipelines.native_transformer import materialize_zero_init
+    return materialize_zero_init
+
+
+def run_zero_init_regression(port):
+    """A finetune predating the last.up/last.down branch omits both keys. Zero-filling them
+    must reproduce the base's dormant behavior: identical output to a model whose up is zeroed
+    (up(down(x)) == 0), while a genuinely missing weight stays a hard mismatch."""
+    materialize_zero_init = load_materialize_zero_init()
+    branch_keys = ("last.up.weight", "last.down.weight")
+
+    torch.manual_seed(0)
+    full_sd = port.Krea2Transformer2DModel(**CFG).state_dict()
+
+    # Ground truth: the branch dormant exactly as the base ships it (up zeroed, down arbitrary).
+    dormant = port.Krea2Transformer2DModel(**CFG).float().eval()
+    dormant_sd = dict(full_sd)
+    dormant_sd["last.up.weight"] = torch.zeros_like(dormant_sd["last.up.weight"])
+    dormant.load_state_dict(dormant_sd)
+
+    # Under test: a checkpoint that omits both branch keys; the loader zero-fills them. A real
+    # missing weight (blocks.0.attn.wq.weight) is dropped too, to confirm it is NOT zero-filled.
+    filled = port.Krea2Transformer2DModel(**CFG).float().eval()
+    hard_key = "blocks.0.attn.wq.weight"
+    partial_sd = {k: v for k, v in full_sd.items() if k not in branch_keys and k != hard_key}
+    missing, unexpected = filled.load_state_dict(partial_sd, strict=False)
+    assert not unexpected, f"unexpected keys: {unexpected}"
+    assert set(missing) == set(branch_keys) | {hard_key}, f"unexpected missing set: {missing}"
+
+    remaining = materialize_zero_init(filled, missing, ("last.down.", "last.up."))
+    assert remaining == [hard_key], f"expected only {hard_key} to remain hard-missing, got {remaining}"
+    for k in branch_keys:
+        w = dict(filled.named_parameters())[k]
+        assert w.abs().sum().item() == 0.0, f"{k} was not zero-filled"
+
+    # Restore the genuinely-missing weight so the forward is well defined, then compare.
+    filled.load_state_dict({hard_key: full_sd[hard_key]}, strict=False)
+
+    img, context, timestep, pos, mask = make_inputs()
+
+    def forward(model):
+        with torch.no_grad():
+            return model(
+                hidden_states=img, encoder_hidden_states=context, timestep=timestep,
+                position_ids=pos, attention_mask=mask, return_dict=False,
+            )[0]
+
+    diff = (forward(dormant) - forward(filled)).abs().max().item()
+    print(f"zero-init regression: max abs diff vs dormant base: {diff:.3e}")
+    assert diff == 0.0, f"ZERO-INIT REGRESSION FAILED: filled output differs by {diff:.3e}"
+    print("ZERO-INIT OK")
+
+
+def main():
+    port = load_port()
+    # Parity runs first in pristine torch state; the regression imports modules afterwards.
+    mmdit = load_reference()
+    run_parity(mmdit, port)
+    run_zero_init_regression(port)
+
+
 if __name__ == "__main__":
     main()