Fix #18: Strict runtime integrity guards for model loading

[Dikshant-Neupane]

Description

This PR addresses Issue #18 by fixing silent failure states during model loading, parameter shape expansion, and context management. It introduces strict runtime guards to ensure that architectural mismatches and missing critical weights loudly fail rather than running silently with garbage initialization.

Key Changes

1. Strict Checkpoint Validation (mage_flow.py, pipeline.py, utils.py)

   • Added CRITICAL_LAYERS definition (img_in, txt_in, proj_out, final_layer.linear).
   • Implemented validate_state_dict_keys to wrap load_state_dict. Checkpoint loading now throws an explicit KeyError with a diagnostic message if a critical layer is missing.

2. Safe Shape Expansion Guards (utils.py)

   • Refactored optionally_expand_state_dict(). If a shape mismatch occurs on a critical projection layer, it now throws a ValueError instead of silently zero-padding (which would produce dead feature channels).
   • Non-critical shape expansions now emit a structured logger.warning.

3. Robust Context Manager Cleanup (mage_text.py)

   • Removed the except Exception: pass anti-pattern in _full_output_mode.
   • The restoration block now uses a try...except wrapper that logs the exact failure via loguru and re-raises the exception to prevent the model from getting stuck in an invalid output mode.

4. VAE-Transformer Mismatch Assertion (pipeline.py)

   • Added an initialization guard that compares the VAE latent_channels against the Transformer in_channels.
   • Triggers an immediate ValueError if the user accidentally pairs mismatched architectures (e.g., a 16-channel VAE with an 8-channel Transformer).

5. Unit Tests Added

   • Introduced tests/test_integrity_guards.py providing a lightweight pytest suite to verify that these specific exceptions trigger reliably without false positives.

Related Issues

• Fixes [BUG] Silent state_dict dropping, unchecked zero-padding, and missing VAE-Transformer channel assertions #18

[Dikshant-Neupane]

@microsoft-github-policy-service agree