Fix #18: Strict runtime integrity guards for model loading

mage_flow/models/mage_flow.py

@@ -174,7 +174,8 @@ def __init__(self, config: ModelConfig):
                 # Handle wrapped EMA format: {'ema_state_dict': ..., ...}
                 if isinstance(sd, dict) and "ema_state_dict" in sd:
                     sd = sd["ema_state_dict"]
-                missing, unexpected = self.load_state_dict(sd, strict=False)
+                from .utils import validate_state_dict_keys
+                missing, unexpected = validate_state_dict_keys(self, sd)
                 if missing:
                     logger.warning(f"Full model load missing keys ({len(missing)}): {missing[:5]}...")
                 if unexpected:

mage_flow/models/modules/mage_text.py

@@ -231,8 +231,10 @@ def _full_output_mode(hf):
                 hf.set_output_mode(prev_mode)
             if prev_skip is not None:
                 hf._skip_lm_head = prev_skip
-        except Exception:  # noqa: BLE001
-            pass
+        except Exception as e:
+            from loguru import logger
+            logger.error(f"Failed to restore text model output mode: {e}")
+            raise
 
 
 def check_edit(model, prompt: str, ref_images, max_new_tokens: int = 192) -> FilterVerdict:

mage_flow/models/utils.py

@@ -12,6 +12,19 @@
 from .mage_flow import MageFlow, MageFlowParams
 
 
+CRITICAL_LAYERS = {"img_in.weight", "txt_in.weight", "proj_out.weight", "final_layer.linear.weight"}
+
+def validate_state_dict_keys(model: torch.nn.Module, state_dict: dict, strict_critical: bool = True):
+    missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
+    if strict_critical:
+        for missing_key in missing:
+            if any(missing_key.endswith(c) for c in CRITICAL_LAYERS):
+                raise KeyError(
+                    f"Critical layer '{missing_key}' is missing from the checkpoint."
+                )
+    return missing, unexpected
+
+
 def get_noise(
     num_samples: int,
     channel: int,
@@ -122,7 +135,7 @@ def load_model_weight(model, pretrain_path, device="cpu"):
 
             sd = correct_model_weight(sd)
             sd = optionally_expand_state_dict(model, sd)
-            missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+            missing, unexpected = validate_state_dict_keys(model, sd)
             print_load_warning(missing, unexpected)
             return True
         except Exception as e:
@@ -159,10 +172,16 @@ def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> di
     """
     Optionally expand the state dict to match the model's parameters shapes.
     """
+    critical_projections = {"img_in", "txt_in", "proj_out"}
     for name, param in model.named_parameters():
         if name in state_dict:
             if state_dict[name].shape != param.shape:
-                logger.info(
+                if any(c in name for c in critical_projections):
+                    raise ValueError(
+                        f"Zero-padding critical projection weights produces dead feature channels. "
+                        f"Cannot safely expand {name} from {state_dict[name].shape} to {param.shape}."
+                    )
+                logger.warning(
                     f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape "
                     f"{param.shape}."
                 )

mage_flow/pipeline.py

@@ -749,7 +749,18 @@ def _resolve(p):
     model = MageFlowModel(cfg)
     sd = load_file(_safe_subpath(repo_dir, "transformer", "diffusion_pytorch_model.safetensors"),
                    device="cpu")
-    model.transformer.load_state_dict(sd, strict=False, assign=True)
+
+    from .models.utils import validate_state_dict_keys
+    validate_state_dict_keys(model.transformer, sd)
+
+    vae_channels = getattr(model.vae, "latent_channels", None)
+    transformer_channels = getattr(model.transformer, "in_channels", None)
+    if vae_channels is not None and transformer_channels is not None:
+        if vae_channels != transformer_channels:
+            raise ValueError(
+                f"Pipeline Architecture Mismatch: VAE latent_channels ({vae_channels}) "
+                f"does not match Transformer in_channels ({transformer_channels})."
+            )
     model.to(device)
     model.transformer.to(torch.bfloat16)
     model.txt_enc.to(torch.bfloat16)

tests/test_integrity_guards.py

@@ -0,0 +1,75 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from mage_flow.models.utils import optionally_expand_state_dict, validate_state_dict_keys
+from mage_flow.models.modules.mage_text import _full_output_mode
+
+
+class MockModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.img_in = nn.Linear(16, 64)
+        self.txt_in = nn.Linear(32, 64)
+        self.proj_out = nn.Linear(64, 16)
+        self.non_critical = nn.Linear(10, 10)
+
+
+def test_validate_state_dict_keys_missing_critical():
+    model = MockModel()
+    state_dict = model.state_dict()
+    # Remove a critical layer
+    del state_dict["img_in.weight"]
+
+    with pytest.raises(KeyError, match="Critical layer 'img_in.weight' is missing"):
+        validate_state_dict_keys(model, state_dict, strict_critical=True)
+
+
+def test_validate_state_dict_keys_success():
+    model = MockModel()
+    state_dict = model.state_dict()
+    # Should not raise
+    missing, unexpected = validate_state_dict_keys(model, state_dict, strict_critical=True)
+    assert len(missing) == 0
+
+
+def test_optionally_expand_state_dict_critical_projection():
+    model = MockModel()
+    state_dict = model.state_dict()
+
+    # Create a shape mismatch on a critical layer
+    state_dict["img_in.weight"] = torch.randn(64, 8) # Expected is 64, 16
+
+    with pytest.raises(ValueError, match="Zero-padding critical projection weights produces dead feature channels"):
+        optionally_expand_state_dict(model, state_dict)
+
+
+def test_optionally_expand_state_dict_non_critical():
+    model = MockModel()
+    state_dict = model.state_dict()
+
+    # Create a shape mismatch on a non-critical layer
+    state_dict["non_critical.weight"] = torch.randn(5, 5) # Expected 10, 10
+
+    # Should not raise and should expand
+    expanded_dict = optionally_expand_state_dict(model, state_dict)
+    assert expanded_dict["non_critical.weight"].shape == (10, 10)
+
+
+def test_full_output_mode_exception_reraise():
+    class MockHF:
+        def __init__(self):
+            self._output_mode = "embedding"
+            self._skip_lm_head = True
+
+        def set_output_mode(self, mode):
+            # simulate an error during restore
+            if mode == "embedding":
+                raise RuntimeError("Simulated failure in set_output_mode")
+            self._output_mode = mode
+
+    hf_model = MockHF()
+
+    with pytest.raises(RuntimeError, match="Simulated failure in set_output_mode"):
+        with _full_output_mode(hf_model):
+            pass # Exception should be raised when exiting the context manager