minimax_h3_nvfp4_checkpoint

def write_nvfp4_decoder_safetensors(path: str, checkpoint: dict[str, Any]) -> None:
    """Write ``{"format", "metadata", "decoder", "post_quant_conv"}`` as one safetensors file."""
    from safetensors.torch import save_file
if checkpoint["format"] not in NVFP4_DECODER_FORMATS:
    raise ValueError(f"unknown NVFP4 decoder format {checkpoint['format']!r}")
tensors = {
    f"{group}.{key}": value.detach().contiguous().cpu().clone()
    for group in GROUPS for key, value in checkpoint[group].items()
}
header = {"format": checkpoint["format"], "nvfp4_metadata": json.dumps(checkpoint["metadata"])}
save_file(tensors, path, metadata=header)</code></pre></details>
fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint.read_nvfp4_decoder_safetensors
read_nvfp4_decoder_safetensors(path: str, device: torch.device | str = 'cpu') -> dict[str, Any]
Read either safetensors format; packed tensors go straight to device (bf16 masters stay on CPU).
Source
def read_nvfp4_decoder_safetensors(path: str, device: torch.device | str = "cpu") -> dict[str, Any]:
    """Read either safetensors format; packed tensors go straight to ``device`` (bf16 masters stay on CPU)."""
    from safetensors import safe_open
    from safetensors.torch import load_file
with safe_open(path, framework=&quot;pt&quot;) as handle:
    header = handle.metadata() or &#123;&#125;
fmt = header.get(&quot;format&quot;)
if fmt not in NVFP4_DECODER_FORMATS:
    raise ValueError(f&quot;&#123;path&#125; is not an NVFP4 decoder safetensors file (format: &#123;fmt!r&#125;); &quot;
                     f&quot;expected one of &#123;NVFP4_DECODER_FORMATS&#125;&quot;)
if &quot;nvfp4_metadata&quot; not in header:
    raise ValueError(f&quot;&#123;path&#125; is missing the nvfp4_metadata header&quot;)
tensors = load_file(path, device=str(device) if fmt == NVFP4_DECODER_PACKED_FORMAT else &quot;cpu&quot;)
checkpoint: dict&#91;str, Any&#93; = &#123;&quot;format&quot;: fmt, &quot;metadata&quot;: json.loads(header&#91;&quot;nvfp4_metadata&quot;&#93;)&#125;
for group in GROUPS:
    prefix = f&quot;&#123;group&#125;.&quot;
    checkpoint&#91;group&#93; = &#123;key&#91;len(prefix):&#93;: value for key, value in tensors.items() if key.startswith(prefix)&#125;
unknown = &#91;key for key in tensors if not key.startswith(tuple(f&quot;&#123;group&#125;.&quot; for group in GROUPS))&#93;
if unknown:
    raise ValueError(f&quot;&#123;path&#125; has tensors outside &#123;GROUPS&#125;: &#123;unknown&#91;:5&#93;&#125;&quot;)
return checkpoint</code></pre></details>
fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint.packed_nvfp4_decoder_checkpoint
packed_nvfp4_decoder_checkpoint(vae: nn.Module, metadata: dict[str, Any]) -> dict[str, Any]
The packed-format checkpoint of a VAE whose NVFP4 decoder linears are all frozen.
Source
def packed_nvfp4_decoder_checkpoint(vae: nn.Module, metadata: dict[str, Any]) -> dict[str, Any]:
    """The packed-format checkpoint of a VAE whose NVFP4 decoder linears are all frozen."""
    decoder = vae.decoder
    state = dict(decoder.state_dict())
    layers = {name: module for name, module in decoder.named_modules() if isinstance(module, NVFP4DecoderLinear)}
    if not layers:
        raise ValueError("the VAE decoder has no NVFP4 linears")
    for name, layer in layers.items():
        if not layer.frozen:
            raise ValueError(f"NVFP4 linear {name} is not frozen; freeze the decoder before exporting packed weights")
        for buffer in PACKED_BUFFERS:
            state[f"{name}.{buffer}"] = getattr(layer, buffer)
        if layer.bias is not None:
            state[f"{name}.bias"] = layer.bias.detach().to(layer.compute_dtype)
    packed_metadata = {**metadata, "packed_layout": PACKED_LAYOUT, "num_linears": len(layers)}
    return {
        "format": NVFP4_DECODER_PACKED_FORMAT,
        "metadata": packed_metadata,
        "decoder": state,
        "post_quant_conv": dict(vae.post_quant_conv.state_dict()),
    }
fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint.apply_packed_nvfp4_decoder
apply_packed_nvfp4_decoder(vae: nn.Module, checkpoint: dict[str, Any]) -> None
Replace the (already depth-cut) dense decoder linears with frozen layers built from packed tensors.
The dense linears are dropped one at a time, so no full-precision copy of a quantized weight is
ever made. Everything else is loaded strictly.
Source
def apply_packed_nvfp4_decoder(vae: nn.Module, checkpoint: dict[str, Any]) -> None:
    """Replace the (already depth-cut) dense decoder linears with frozen layers built from packed tensors.
The dense linears are dropped one at a time, so no full-precision copy of a quantized weight is
ever made. Everything else is loaded strictly.
&quot;&quot;&quot;
metadata = checkpoint&#91;&quot;metadata&quot;&#93;
decoder = vae.decoder
state = dict(checkpoint&#91;&quot;decoder&quot;&#93;)
names = nvfp4_decoder_linear_names(decoder, tuple(metadata.get(&quot;skip_blocks&quot;) or ()))
check_num_linears(metadata, len(names))
act_scale = metadata.get(&quot;act_scale&quot;, &quot;dynamic&quot;)
for name in names:
    if act_scale != &quot;static&quot;:
        # Only static scaling reads input_amax, so a file may omit it; the strict load still needs the key.
        state.setdefault(f&quot;&#123;name&#125;.input_amax&quot;, torch.zeros((), dtype=torch.float32))
    required = PACKED_BUFFERS + (&quot;input_amax&quot;, )
    missing = &#91;f&quot;&#123;name&#125;.&#123;buffer&#125;&quot; for buffer in required if f&quot;&#123;name&#125;.&#123;buffer&#125;&quot; not in state&#93;
    if missing:
        raise ValueError(f&quot;the packed NVFP4 decoder does not fit this VAE&#x27;s architecture: missing &#123;missing&#125;&quot;)
    parent_name, _, child = name.rpartition(&quot;.&quot;)
    parent = decoder.get_submodule(parent_name)
    dense = parent&#91;int(child)&#93; if child.isdigit() else getattr(parent, child)
    packed = state.pop(f&quot;&#123;name&#125;.packed_weight&quot;)
    if packed.shape&#91;0&#93; != dense.out_features:
        raise ValueError(f&quot;the packed NVFP4 decoder does not fit this VAE&#x27;s architecture: &#123;name&#125; has &quot;
                         f&quot;&#123;packed.shape&#91;0&#93;&#125; outputs, the VAE &#123;dense.out_features&#125;&quot;)
    layer = NVFP4DecoderLinear.from_packed(packed,
                                           state.pop(f&quot;&#123;name&#125;.weight_inv_scale&quot;),
                                           state.pop(f&quot;&#123;name&#125;.weight_global_sf&quot;),
                                           state.get(f&quot;&#123;name&#125;.bias&quot;),
                                           state&#91;f&quot;&#123;name&#125;.input_amax&quot;&#93;,
                                           in_features=dense.in_features,
                                           rotation_group=metadata.get(&quot;rotation_group&quot;),
                                           compute_dtype=torch.bfloat16,
                                           act_scale=act_scale)
    if child.isdigit():
        parent&#91;int(child)&#93; = layer
    else:
        setattr(parent, child, layer)
    del dense
enable_fused_blocks(decoder)
try:
    decoder.load_state_dict(state, strict=True)
    vae.post_quant_conv.load_state_dict(checkpoint&#91;&quot;post_quant_conv&quot;&#93;, strict=True)
except RuntimeError as error:
    raise ValueError(f&quot;the packed NVFP4 decoder does not fit this VAE&#x27;s architecture: &#123;error&#125;&quot;) from error
vae.requires_grad_(False)</code></pre></details>