Skip to content

[bugfix]: read H3 geometry from checkpoint JSON so lazy load can drop the DiT before VAE decode - #2

Merged
Satyam-53 merged 2 commits into
KyleNeverGivesUp:lazy-module-loadfrom
aryan5v:aryan/1761-lazy-geometry
Aug 31, 2026
Merged

[bugfix]: read H3 geometry from checkpoint JSON so lazy load can drop the DiT before VAE decode#2
Satyam-53 merged 2 commits into
KyleNeverGivesUp:lazy-module-loadfrom
aryan5v:aryan/1761-lazy-geometry

Conversation

@aryan5v

@aryan5v aryan5v commented Aug 31, 2026

Copy link
Copy Markdown

Summary

  • Input prep and video decode no longer hold live VAE/DiT proxies just to read spatial_compression_ratio, latent_channels, and patch_size. Those scalars now come from checkpoint config.json via update_model_arch (no weights).
  • Last-holder scheduling can therefore drop Qwen after conditioning, drop the DiT after denoise, and load the VAE only for decode. Spark alpine 768×1344×124 with --lazy-module-load confirmed that order (VAE 10.34 GiB after DiT release of 65.51 GiB).
  • --lazy-module-load is auto-on for unified memory (None default) and for basic_fasth3.py / basic_minimax_h3_t2v.py when --num-gpus 1. --no-lazy-module-load still keeps every component resident. Docs now describe the opt-in flag instead of always-on sequential start.

This is the review follow-up for hao-ai-lab/FastVideo#1761. It does not include CUDA TAEH3.

Test plan

… the DiT before VAE decode

Input prep and unpatchify were holding live VAE/DiT proxies just for two integers, which loaded the video VAE before Qwen and kept the DiT resident through decode. Auto-enable --lazy-module-load on unified memory and on single-GPU FastH3 examples.
…ing docs

The MLX phase-order sentence and "start a new generator" instruction belong
to a different design. A later generate() reloads in-process.
@Satyam-53

Copy link
Copy Markdown

/merge

@Satyam-53
Satyam-53 merged commit 770d2df into KyleNeverGivesUp:lazy-module-load Aug 31, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants