Fix/issue 19

[Dikshant-Neupane]

Description

This PR addresses Issue #19, refactoring the Mage-VAE latent tokenizer (mage_vae.py and utils.py) to eliminate spatial artifacts, prevent precision quantization noise, and enforce deterministic local RNG.

Key Changes

1. Reflection Padding Guard (utils.py, mage_vae.py)

   • Added a pad_to_patch_multiple helper that applies dynamic spatial reflection padding to ensure input images are exact multiples of 16.
   • Replaced fragile math.ceil division with safe integer division //.
   • Prevents boundary smearing and sub-pixel edge repetition on extreme aspect ratios.

2. FP32 Precision Guard for Posterior Sampling (mage_vae.py)

   • Upcasts mean and logvar to torch.float32 inside encode() prior to computing torch.exp(0.5 * logvar).
   • Generates noise and scales standard deviation strictly in FP32, eliminating the 12.5% relative quantization noise caused by 7-bit mantissa truncation in bfloat16.

3. Dynamic adaLN Cache Management (mage_vae.py)

   • Refactored _ConstAdaLN to retain a reference to the original_mlp alongside the constant-folded buffer.
   • The decoder now dynamically inspects t; if a non-zero timestep is detected (t > 1e-5), the _ConstAdaLN modules gracefully bypass the cached buffer and execute the live MLP, properly supporting multi-step progressive decoding paths.
   • Added set_adaln_cache() method to MageVAE for explicit cache state control.

4. Deterministic Local RNG (mage_vae.py)

   • Modified the encode() signature to accept an explicit generator: torch.Generator | None = None.
   • Replaced torch.randn_like with torch.randn(..., generator=generator). This decouples the sampling process from PyTorch's volatile global RNG state, guaranteeing deterministic latents across runs.

5. Test Suite (tests/test_vae_integrity.py)

   • Implemented a complete test suite validating exact patch multiples, strict generation determinism under fixed seeds, successful bypass routing for dynamic adaLN caches, and the absence of NaN/underflow errors during FP32/BF16 boundary scaling.

Related Issues

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