diff --git a/mage_flow/models/mage_flow.py b/mage_flow/models/mage_flow.py
index 69dabe3..c1b0f33 100644
--- a/mage_flow/models/mage_flow.py
+++ b/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:
diff --git a/mage_flow/models/modules/mage_text.py b/mage_flow/models/modules/mage_text.py
index 5255262..58c74b5 100644
--- a/mage_flow/models/modules/mage_text.py
+++ b/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:
diff --git a/mage_flow/models/modules/mage_vae.py b/mage_flow/models/modules/mage_vae.py
index 58c1b10..3ecb91d 100644
--- a/mage_flow/models/modules/mage_vae.py
+++ b/mage_flow/models/modules/mage_vae.py
@@ -341,15 +341,20 @@ def to_patches(t):
 # torch.compile can fuse the surrounding ops normally.
 # ---------------------------------------------------------------------------
 class _ConstAdaLN(nn.Module):
-    def __init__(self, modulation: torch.Tensor):
+    def __init__(self, modulation: torch.Tensor, original_mlp: nn.Module):
         super().__init__()
         self.register_buffer("modulation", modulation.detach().clone())
+        self.original_mlp = original_mlp
+        self.enabled = True
+        self.bypass = False
 
     def forward(self, c):
-        b = c.shape[0]
-        if self.modulation.shape[0] != b:
-            return self.modulation.expand(b, *self.modulation.shape[1:])
-        return self.modulation
+        if getattr(self, "enabled", True) and not getattr(self, "bypass", False):
+            b = c.shape[0]
+            if self.modulation.shape[0] != b:
+                return self.modulation.expand(b, *self.modulation.shape[1:])
+            return self.modulation
+        return self.original_mlp(c)
 
 
 def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
@@ -365,7 +370,7 @@ def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
             continue
         with torch.no_grad():
             mod = adaln(c)
-        child.adaLN_modulation = _ConstAdaLN(mod)
+        child.adaLN_modulation = _ConstAdaLN(mod, original_mlp=adaln)
         n += 1
     return n
 
@@ -491,6 +496,11 @@ def __init__(
     def forward(self, x, t, cond):
         b, _, h, w = x.shape
         c = self.t_embedder(t.view(-1))
+        
+        bypass_cache = t.abs().max().item() > 1e-5
+        for m in self.blocks.modules():
+            if isinstance(m, _ConstAdaLN):
+                m.bypass = bypass_cache
 
         s = self.s_embedder(x, cond)
         for block in self.blocks:
@@ -544,6 +554,7 @@ class MageVAE(nn.Module):
     def __init__(self, ckpt_path: str, sample_posterior: bool = True):
         super().__init__()
         self.sample_posterior = sample_posterior
+        self._adaln_cache_enabled = True
 
         self.dconv_encoder = _DConvEncoder()
         self.decoder_model = _DConvDenoiser()
@@ -612,14 +623,20 @@ def _encode_moments(self, x: torch.Tensor):
         return self._moments(x)
 
     @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> torch.Tensor:
+    def encode(self, x: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        from ..utils import pad_to_patch_multiple
+        
         ps = self.dconv_encoder.patch_size
-        H, W = x.shape[-2], x.shape[-1]
-        if H % ps or W % ps:
-            raise ValueError(f"H, W must be multiples of {ps}, got ({H}, {W})")
+        x, (pad_h, pad_w) = pad_to_patch_multiple(x, patch_size=ps)
+        
         mean, logvar = self._encode_moments(x)
         if self.sample_posterior:
-            return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean)
+            mean_f32 = mean.float()
+            logvar_f32 = logvar.float()
+            std_f32 = torch.exp(0.5 * logvar_f32)
+            eps = torch.randn(mean_f32.shape, device=mean.device, dtype=torch.float32, generator=generator)
+            latent_f32 = mean_f32 + std_f32 * eps
+            return latent_f32.to(dtype=mean.dtype)
         return mean
 
     @torch.no_grad()
@@ -649,3 +666,10 @@ def _freeze_adaln_cache(self):
         _replace_adaln_with_const(self.dconv_encoder, c_enc)
         c_dec = self.decoder_model.t_embedder(t)
         _replace_adaln_with_const(self.decoder_model, c_dec)
+        self.set_adaln_cache(self._adaln_cache_enabled)
+
+    def set_adaln_cache(self, enabled: bool):
+        self._adaln_cache_enabled = enabled
+        for m in self.modules():
+            if isinstance(m, _ConstAdaLN):
+                m.enabled = enabled
diff --git a/mage_flow/models/utils.py b/mage_flow/models/utils.py
index 30bb3a2..c648c04 100644
--- a/mage_flow/models/utils.py
+++ b/mage_flow/models/utils.py
@@ -8,10 +8,34 @@
 from safetensors.torch import load_file
 from safetensors.torch import load_file as load_sft
 from torch import Tensor
+from typing import Tuple
+import torch.nn.functional as F
 
 from .mage_flow import MageFlow, MageFlowParams
 
 
+def pad_to_patch_multiple(x: torch.Tensor, patch_size: int = 16) -> Tuple[torch.Tensor, Tuple[int, int]]:
+    H, W = x.shape[-2], x.shape[-1]
+    pad_h = (patch_size - H % patch_size) % patch_size
+    pad_w = (patch_size - W % patch_size) % patch_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
+    return x, (pad_h, pad_w)
+
+
+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,
@@ -25,8 +49,8 @@ def get_noise(
     return torch.randn(
         num_samples,
         channel,
-        math.ceil(height / 16),
-        math.ceil(width / 16),
+        height // 16,
+        width // 16,
         device=device,
         dtype=dtype,
         generator=torch.Generator(device=device).manual_seed(seed),
@@ -38,8 +62,8 @@ def unpack(x: Tensor, height: int, width: int) -> Tensor:
     return rearrange(
         x,
         "b (h w) c -> b c h w",
-        h=math.ceil(height / 16),
-        w=math.ceil(width / 16),
+        h=height // 16,
+        w=width // 16,
     )
 
 
@@ -122,7 +146,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 +183,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}."
                 )
diff --git a/mage_flow/pipeline.py b/mage_flow/pipeline.py
index 3e5aaba..a3bd3d8 100644
--- a/mage_flow/pipeline.py
+++ b/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)
diff --git a/tests/test_integrity_guards.py b/tests/test_integrity_guards.py
new file mode 100644
index 0000000..74092e1
--- /dev/null
+++ b/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
diff --git a/tests/test_vae_integrity.py b/tests/test_vae_integrity.py
new file mode 100644
index 0000000..7ae7b96
--- /dev/null
+++ b/tests/test_vae_integrity.py
@@ -0,0 +1,106 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from mage_flow.models.modules.mage_vae import _ConstAdaLN
+from mage_flow.models.utils import pad_to_patch_multiple
+
+
+def test_pad_to_patch_multiple():
+    # Test tensor of shape [1, 3, 50, 50] padded to multiple of 16
+    x = torch.randn(1, 3, 50, 50)
+    x_pad, (pad_h, pad_w) = pad_to_patch_multiple(x, patch_size=16)
+    
+    assert x_pad.shape[-2] == 64
+    assert x_pad.shape[-1] == 64
+    assert pad_h == 14
+    assert pad_w == 14
+    
+    # Test already multiple
+    y = torch.randn(1, 3, 32, 32)
+    y_pad, (pad_h, pad_w) = pad_to_patch_multiple(y, patch_size=16)
+    assert y_pad.shape[-2] == 32
+    assert y_pad.shape[-1] == 32
+    assert pad_h == 0
+    assert pad_w == 0
+
+
+def test_deterministic_encoding():
+    class MockVAE(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.sample_posterior = True
+            
+            class MockEncoder:
+                patch_size = 16
+            self.dconv_encoder = MockEncoder()
+            
+        def _encode_moments(self, x):
+            B, C, H, W = x.shape
+            mean = torch.zeros(B, 128, H // 16, W // 16)
+            logvar = torch.zeros(B, 128, H // 16, W // 16)
+            return mean, logvar
+            
+    vae = MockVAE()
+    from mage_flow.models.modules.mage_vae import MageVAE as ActualMageVAE
+    vae.encode = ActualMageVAE.encode.__get__(vae)
+    
+    x = torch.randn(1, 3, 64, 64)
+    
+    gen1 = torch.Generator().manual_seed(42)
+    out1 = vae.encode(x, generator=gen1)
+    
+    gen2 = torch.Generator().manual_seed(42)
+    out2 = vae.encode(x, generator=gen2)
+    
+    gen3 = torch.Generator().manual_seed(43)
+    out3 = vae.encode(x, generator=gen3)
+    
+    assert torch.allclose(out1, out2)
+    assert not torch.allclose(out1, out3)
+
+
+def test_fp32_precision_guard():
+    class MockVAE(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.sample_posterior = True
+            class MockEncoder:
+                patch_size = 16
+            self.dconv_encoder = MockEncoder()
+            
+        def _encode_moments(self, x):
+            mean = torch.zeros(1, 128, 4, 4, dtype=torch.bfloat16)
+            logvar = torch.full((1, 128, 4, 4), -20.0, dtype=torch.bfloat16)
+            return mean, logvar
+            
+    vae = MockVAE()
+    from mage_flow.models.modules.mage_vae import MageVAE as ActualMageVAE
+    vae.encode = ActualMageVAE.encode.__get__(vae)
+    
+    x = torch.randn(1, 3, 64, 64, dtype=torch.bfloat16)
+    
+    gen = torch.Generator().manual_seed(0)
+    out = vae.encode(x, generator=gen)
+    
+    assert not torch.isnan(out).any()
+    assert out.dtype == torch.bfloat16
+
+
+def test_dynamic_adaln_cache():
+    original_mlp = nn.Linear(4, 4)
+    nn.init.constant_(original_mlp.weight, 1.0)
+    nn.init.constant_(original_mlp.bias, 1.0)
+    modulation = torch.zeros(1, 4)
+    
+    const_adaln = _ConstAdaLN(modulation, original_mlp)
+    c_t0 = torch.zeros(1, 4)
+    
+    # Enabled and no bypass -> returns modulation
+    const_adaln.enabled = True
+    const_adaln.bypass = False
+    assert torch.allclose(const_adaln(c_t0), modulation)
+    
+    # Bypass -> returns original_mlp output
+    const_adaln.bypass = True
+    assert not torch.allclose(const_adaln(c_t0), modulation)