Diffusers documentation
AutoencoderKLMiniMaxH3Image
AutoencoderKLMiniMaxH3Image
The image variational autoencoder (VAE) model with KL loss used in MiniMax-H3-Image by MiniMax. It pairs the causal 3D CNN encoder of AutoencoderKLMiniMaxH3, run over a single frame, with a 2D ViT decoder, and compresses 16x spatially. encode and decode take and return 4D image tensors.
It shares three conventions with AutoencoderKLMiniMaxH3:
- Latents are normalized per channel. There is no
scaling_factor: a pipeline encodes with(latent - latents_mean) / latents_stdand decodes withlatent * latents_std + latents_mean. - The pixel convention is ImageNet-normalized RGB over a
[0, 1]base range, not the usual[-1, 1].encodeexpects(pixel - imagenet_mean) / imagenet_stdanddecodereturns values in that same space, so a pipeline appliessample * imagenet_std + imagenet_meanand clamps to[0, 1]before postprocessing. - Spatial tiling is on by default. MiniMax-H3-Image was released with tiling enabled for both encoding and decoding and the released images are the blended-tile ones, so turning it off changes the output. Use
enable_tilingto change the tile geometry anddisable_tilingto switch it off.
import torch
from diffusers import AutoencoderKLMiniMaxH3Image
vae = AutoencoderKLMiniMaxH3Image.from_pretrained(
"MiniMaxAI/MiniMax-H3-Image", subfolder="vae_diffusers", dtype=torch.float32
).to("cuda") # or "mps", "xpu", "cpu"AutoencoderKLMiniMaxH3Image
class diffusers.AutoencoderKLMiniMaxH3Image
< source >( in_channels: int = 3out_channels: int = 3latent_channels: int = 24block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024)layers_per_block: int = 2spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1)temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1)norm_num_groups: int = 32norm_eps: float = 1e-06spatial_padding_mode: str = 'reflect'decoder_num_layers: int = 24decoder_num_attention_heads: int = 16decoder_attention_head_dim: int = 64decoder_num_register_tokens: int = 4decoder_ffn_mult: int = 4decoder_rope_theta: float = 100.0decoder_rope_dim_ratio: float = 0.5decoder_norm_eps: float = 1e-05latents_mean: tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)latents_std: tuple[float, ...] = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0) )
A VAE model with a causal 3D CNN encoder and a non-causal 2D ViT decoder, used in MiniMax-H3-Image. The encoder is the MiniMax-H3 video encoder run over a single frame; the decoder is an image-only ViT.
This model inherits from ModelMixin. Check the superclass documentation for it’s generic methods implemented for all models (such as downloading or saving).
Latents are normalized with per-channel latents_mean / latents_std rather than a scaling_factor; a pipeline
encodes with (latent - latents_mean) / latents_std and decodes with latent * latents_std + latents_mean.
The pixel convention is ImageNet-normalized RGB over a [0, 1] base range, not the usual [-1, 1]: encode expects (pixel - imagenet_mean) / imagenet_std and decode returns values in that same space, so a pipeline has
to apply sample * imagenet_std + imagenet_mean (mean (0.485, 0.456, 0.406), std (0.229, 0.224, 0.225)) and
clamp to [0, 1] before postprocessing.
Unlike most autoencoders in the library, spatial tiling is on by default: MiniMax-H3-Image was released with
tiling enabled for both encoding and decoding, and the released images are the blended-tile ones, so disabling
tiling changes the output. Use enable_tiling to change the tile geometry, disable_tiling to turn it off.
encode
< source >( x: Tensorreturn_dict: bool = True )
Parameters
- x (
torch.Tensor) — Input batch of images, shape(batch_size, in_channels, height, width). - return_dict (
bool, optional, defaults toTrue) — Whether to return a AutoencoderKLOutput instead of a plain tuple.
Encode a batch of images into latents.
decode
< source >( z: Tensorreturn_dict: bool = True ) → DecoderOutput or tuple
Parameters
- z (
torch.Tensor) — Input batch of latents, shape(batch_size, latent_channels, height, width). - return_dict (
bool, optional, defaults toTrue) — Whether to return a DecoderOutput instead of a plain tuple.
Returns
DecoderOutput or tuple
The decoded images, shape (batch_size, out_channels, height * spatial_compression_ratio, width * spatial_compression_ratio).
Decode a batch of latents.
enable_tiling
< source >( tile_sample_min_height: int | None = Nonetile_sample_min_width: int | None = Nonetile_sample_min_overlap_height: int | None = Nonetile_sample_min_overlap_width: int | None = None )
Parameters
- tile_sample_min_height (
int, optional) — The tile height in pixel space. Frames taller than this are split along the height dimension. - tile_sample_min_width (
int, optional) — The tile width in pixel space. Frames wider than this are split along the width dimension. - tile_sample_min_overlap_height (
int, optional) — The minimum overlap, in pixels, between two consecutive vertical tiles. - tile_sample_min_overlap_width (
int, optional) — The minimum overlap, in pixels, between two consecutive horizontal tiles.
Enable tiled VAE encoding/decoding. When this option is enabled, the VAE splits the frames into tiles, encodes or decodes each tile separately and linearly blends the overlaps back together. This lowers the memory requirement and allows processing larger frames.
forward
< source >( sample: Tensorsample_posterior: bool = Falsegenerator: typing.Optional[torch.Generator] = Nonereturn_dict: bool = True ) → DecoderOutput or tuple
Parameters
- sample (
torch.Tensor) — Input batch of images, shape(batch_size, in_channels, height, width). - sample_posterior (
bool, optional, defaults toFalse) — Whether to sample the posterior instead of taking its mode. - generator (
torch.Generator, optional) — Generator used whensample_posterior=True. - return_dict (
bool, optional, defaults toTrue) — Whether to return a DecoderOutput instead of a plain tuple.
Returns
DecoderOutput or tuple
The round-tripped images, shape (batch_size, out_channels, height, width).
Encode then decode a batch of images.