[BUG/VAE] Boundary smearing at extreme aspect ratios, BF16 posterior underflow, and non-deterministic encoding

Description

While reviewing the Mage-VAE tokenizer and latent space sampling implementations (mage_vae.py and utils.py), I identified four distinct bugs affecting image encoding quality, numerical stability, and deterministic execution.

Below is a detailed breakdown of the root causes, reproduction steps, and recommended fixes.

Affected Components & Root Cause Analysis

1. Extreme Aspect Ratio Boundary Artifacts via Non-Replicating Padding

• Locations: utils.py (get_noise) and mage_vae.py (encode)
• Issue: get_noise() uses math.ceil(height / 16) to compute latent shapes, while vae.encode() requires inputs to be strict multiples of 16. At extreme aspect ratios (e.g., $4:1 \to 512 \times 2048$), latent feature maps develop highly anisotropic receptive fields. In AttnBlock, padding with mode="replicate" along the short axis causes boundary features to bleed axially, producing sub-pixel smearing.
• Impact: Distorted edges and color bleeding along image boundaries when working with panoramic or non-standard aspect ratios.

2. BF16 / FP16 Posterior Sampling Quantization Noise

• Location: mage_vae.py (encode)
• Issue: logvar is clamped to [-20, 10]. At the lower bound (logvar = -20), the sampling equation evaluates $\exp(0.5 \times -20) = \exp(-10) \approx 4.54 \times 10^{-5}$. In bfloat16 (7-bit mantissa), evaluation of $\exp(-10)$ retains only ~3 significant bits—introducing ~12.5% quantization noise in low-variance latent channels.
• Impact: Visible color banding and spatial noise artifacts in smooth or uniform image regions when running in bfloat16.

3. Frozen adaLN Cache Blocks Dynamic Timestep Decoding

• Location: mage_vae.py (_freeze_adaln_cache)
• Issue: _freeze_adaln_cache() permanently replaces adaLN_modulation MLPs with constant buffers evaluated strictly at $t=0$.
• Impact: While correct for standard single-pass $t=0$ decoding, any downstream extension relying on iterative VAE refinement or multi-step latent optimization will silently produce incorrect outputs due to the hardcoded modulation cache.

4. Non-Deterministic VAE Posterior Sampling via Global RNG

• Location: mage_vae.py (encode)
• Issue: encode() calls torch.randn_like(mean) using PyTorch's global RNG state instead of accepting a local torch.Generator.
• Impact: Even when a pipeline seed is explicitly set, VAE encoding during edit/refinement steps can draw different random numbers, breaking generation reproducibility across identical seeds.

Suggested Fixes

1. Replace math.ceil with explicit reflection padding in input space:

   def _pad_to_multiple(x: torch.Tensor, ps: int = 16):
       H, W = x.shape[-2], x.shape[-1]
       pad_h = (ps - H % ps) % ps
       pad_w = (ps - W % ps) % ps
       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)

Start by reading get_noise in utils.py and encode, AttnBlock, and _freeze_adaln_cache in mage_vae.py, then reproduce the four reported cases, especially extreme aspect ratios and reduced-precision encoding. The issue is done when boundary artifacts, posterior stability, timestep behavior, and seeded encoding are each verified against the intended behavior.