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

[Dikshant-Neupane]

Overview

During a technical audit of the codebase, I identified several critical bugs, architectural edge cases, and performance bottlenecks across the VAE tokenizer, Rectified Flow sampling, CUDA/kernel backends, and pipeline loading logic.

Below is a detailed breakdown of the root causes, reproduction conditions, and recommended fixes for these issues.

---

1. Latent Tokenizer & Spatial Compression (Mage-VAE)

• 1.1 Extreme Aspect Ratio Boundary Artifacts via Non-Replicating Padding (mage_vae.py, utils.py)

  • Issue: get_noise() uses math.ceil(height/16) while encode() requires exact multiples of 16. At extreme ratios (e.g., 4:1), AttnBlock uses mode="replicate" padding along the short axis, propagating boundary features axially and causing sub-pixel smearing.
  • Fix: Replace math.ceil with explicit reflection padding in input space (F.pad(x, ..., mode="reflect")).

• 1.2 FP16 / BF16 Posterior Sampling Quantization Noise (mage_vae.py)

  • Issue: Lower-bounded logvar = -20 yields $\exp(-10) \approx 4.54 \times 10^{-5}$. In BF16, this has only ~3 significant mantissa bits, introducing ~12.5% quantization noise in uniform image regions.
  • Fix: Compute posterior sampling exponentially in torch.float32 regardless of model dtype before casting back.

• 1.3 Frozen adaLN Modulation Prevents Dynamic Timestep Usage (mage_vae.py)

  • Issue: _freeze_adaln_cache() permanently replaces modulation MLPs with constant buffers at $t=0$. Any progressive or iterative VAE refinement extensions silently fail.
  • Fix: Guard the cache replacement behind an explicit _enable_adaln_cache toggle.

---

2. Flow Matching & Sampling Trajectories (Mage-Flow)

• 2.1 CFG Renormalization Directional Collapse (pipeline.py)

  • Issue: Dividing by per-token norm rescales velocity vectors non-uniformly and gates near-zero tokens (e.g., solid backgrounds) to zero velocity, causing spatial distortion.
  • Fix: Apply global norm rescaling across tokens to preserve flow-field vector angles (cond_norm / (comb_norm + 1e-6)).

• 2.2 Trajectory Drift in Few-Step Turbo Schedules (pipeline.py)

  • Issue: Linear base sigmas with static shift $6.0$ cause a steep 33% step drop from $\sigma=1.0 \to 0.667$ on step 1, where ODE curvature is highest.
  • Fix: Use a cosine distribution for sigmas when num_steps <= 8.

• 2.3 Multi-Image Sequence Packing VRAM Explosion (pipeline.py)

  • Issue: batch_cfg=True duplicates joint token sequences (e.g., $16,384$ tokens for dual $1024 \times 1024$ images), triggering $O(N^2)$ flash-attention workspace allocation OOMs on 24GB GPUs.
  • Fix: Auto-disable batch_cfg when estimated memory exceeds 70–80% of available VRAM.

---

3. CUDA Kernels & Hardware Compatibility

• 3.1 SDPA Loop Overhead on Non-FlashAttention Hardware (_attn_backend.py)

  • Issue: Fallback SDPA paths dispatch Python loops over packed sequences ($4 \times 28 \times 30 = 3,360$ separate kernel launches per generation), causing massive CPU-bound latency bottlenecks on non-A100 hardware.
  • Fix: Write a fused Triton kernel for packing scatter + attention or batch the sequence slice calls.

• 3.2 Silent FA4 Fallback / Crashes on Ampere GPUs (_attn_backend.py)

  • Issue: Window size normalization converts (-1, -1) to (None, None) for FlashAttention-4 without checking CUDA capability, causing silent execution fallbacks or opaque CUTE errors on SM80/86 (A100/RTX 3090).
  • Fix: Add an explicit architecture capability check (torch.cuda.get_device_capability() >= (9, 0)).

---

4. Generation & Multimodal Editing Quality Failures

• 4.1 Content Refusal Collision with White Generations (pipeline.py)

  • Issue: make_refusal_image() outputs a plain white RGB image, making policy blocks indistinguishable from valid generations or NaN latent overflows.
  • Fix: Add a subtle visual indicator pattern or watermark to refusal returns.

• 4.2 Multimodal Aspect Ratio Mismatch in Reference Editing (pipeline.py)

  • Issue: Resizing reference images directly to output dimensions squashes non-square ratios (e.g., $1920 \times 1080 \to 1024 \times 1024$), creating a spatial aspect ratio mismatch between VAE latents and text conditioning.
  • Fix: Reflection-pad reference images to the target aspect ratio prior to encoding.

---

5. Codebase & Pipeline Integration Guards

• 5.1 Silent Checkpoint Key Dropping (mage_flow.py, pipeline.py)

  • Issue: load_state_dict(strict=False, assign=True) suppresses missing state dict keys, allowing execution with uninitialized layers when architecture parameters mismatch.
  • Fix: Add strict verification for required model keys upon initialization.

• 5.2 Unchecked Zero-Padding on Parameter Shape Mismatches (utils.py)

  • Issue: optionally_expand_state_dict silently zero-pads mismatched parameter shapes into target tensors, leaving critical layers (e.g., img_in.weight) partially filled with zeros.
  • Fix: Raise explicit errors on shape mismatches for critical projection parameters.

• 5.3 Unhandled Mode Restoration Exceptions (mage_text.py)

  • Issue: Silently caught exceptions in _full_output_mode can leave the model stuck in the wrong mode, causing content filter evaluations to fail.
  • Fix: Explicitly log exceptions and re-throw or fallback safely.

• 5.4 Non-Deterministic VAE Posterior Sampling (mage_vae.py)

  • Issue: encode() uses global torch.randn_like instead of accepting an explicit torch.Generator, causing seed non-reproducibility across runs.
  • Fix: Pass an optional local generator parameter to encode().

• 5.5 Missing VAE Latent vs. Transformer Channel Assertion (pipeline.py)

  • Issue: No runtime check confirms vae.latent_channels == transformer.in_channels, risking silent dimensional slicing on configuration errors.
  • Fix: Add an explicit channel equality check in compute_vae_encodings.

---

Conclusion

I have tested local refactors for these issues and would be glad to submit PRs for any specific subsystems if the maintainers agree with the proposed fixes!