[diffusion] feat: support streaming native vae weights directly to gpu (#37004)

This commit is contained in:
Mick
2026-08-30 20:48:09 +08:00
committed by GitHub
parent e51a3ae65e
commit aa483ab782
9 changed files with 510 additions and 27 deletions
+9
View File
@@ -131,6 +131,7 @@ pipeline's registered module name:
| --- | --- | --- | --- |
| Replace a component | `--component-paths.<component> {MODEL}` | `--<component>-path {MODEL}` | Load the replacement component's configuration and weights |
| Replace only its weights | `--component-weights-paths.<component> {WEIGHTS}` | `--<component>-weights-path {WEIGHTS}` | Retain the base component configuration and replace its weights |
| Direct-load an eligible component | `--component-direct-gpu-weight-loading.<component>` | None | Use that component's audited direct-GPU loader; it must stay resident |
| Quantize an unquantized component online | `--component-quantizations.<component> {METHOD}` | `--<component>-quantization {METHOD}` | Apply a method supported by that component's native loader |
| Keep selected component layers unquantized | `--component-quantization-ignored-layers.<component> {PATTERN...}` | None | Pass component-local ignored-layer patterns to its online quantizer |
@@ -166,6 +167,14 @@ published, model-specific checkpoint examples; for example, all H3 sources and
their exact overlays are kept in one
[MiniMax-H3 compatibility table](/cookbook/diffusion/MiniMax/MiniMax-H3#checkpoint-and-adapter-formats).
Direct-GPU loading is also capability-based. The existing
`--direct-gpu-weight-loading` remains the primary DiT path. The component form
currently supports standard native `vae` and `video_vae` state dicts on CUDA:
it streams each safetensors tensor directly to the resident
module, rather than materializing a complete CPU state dict. It rejects custom
Diffusers `auto_map` classes, quantized checkpoints, tied state entries, and
any component-offload or layerwise-offload placement.
Component overrides accept a local component directory, a standalone Hub
repository, or a Hub component subfolder written as `owner/repo/subfolder`.
For transformer and native encoder loaders, an explicit weight filename keeps