Made it possible to use the flux 2 small decoder VAE when setting the vae_path manually for flux2 models

extensions_built_in/diffusion_models/flux2/flux2_model.py

@@ -21,7 +21,7 @@ from toolkit.util.quantize import quantize, get_qtype, quantize_model
 from transformers import AutoProcessor, Mistral3ForConditionalGeneration
 from .src.model import Flux2, Flux2Params
 from .src.pipeline import Flux2Pipeline
-from .src.autoencoder import AutoEncoder, AutoEncoderParams
+from .src.autoencoder import AutoEncoder, AutoEncoderParams, AutoEncoderSmallDecoderParams
 from safetensors.torch import load_file, save_file
 from PIL import Image
 import torch.nn.functional as F
@@ -193,18 +193,30 @@ class Flux2Model(BaseModel):
             vae_path = self.flux2_vae_path
 
         if vae_path is None or not os.path.exists(vae_path):
+            vae_filename = FLUX2_VAE_FILENAME
+            if vae_path is not None:
+                # see if it is a filename for huggingface hub
+                if len(vae_path.split("/")) == 3 and vae_path.endswith(".safetensors"):
+                    vae_filename = vae_path.split("/")[-1]
+                    vae_path = "/".join(vae_path.split("/")[:-1])
             p = vae_path if vae_path is not None else model_path
             # assume it is from the hub
             vae_path = huggingface_hub.hf_hub_download(
                 repo_id=p,
-                filename=FLUX2_VAE_FILENAME,
+                filename=vae_filename,
                 token=HF_TOKEN,
             )
-        with torch.device("meta"):
-            vae = AutoEncoder(AutoEncoderParams())
         
         vae_state_dict = load_file(vae_path, device="cpu")
         
+        autoencoder_params = AutoEncoderParams()
+        if vae_state_dict['decoder.up.0.block.0.conv1.bias'].shape[0] == 96:
+            # this is the small decoder version
+            autoencoder_params = AutoEncoderSmallDecoderParams()
+        
+        with torch.device("meta"):
+            vae = AutoEncoder(autoencoder_params)
+
         # cast to dtype
         for key in vae_state_dict:
             vae_state_dict[key] = vae_state_dict[key].to(dtype)

extensions_built_in/diffusion_models/flux2/src/autoencoder.py

@@ -17,6 +17,17 @@ class AutoEncoderParams:
     num_res_blocks: int = 2
     z_channels: int = 32
 
+@dataclass
+class AutoEncoderSmallDecoderParams:
+    resolution: int = 256
+    in_channels: int = 3
+    ch: int = 128
+    ch_encoder: int = 96
+    out_ch: int = 3
+    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
+    num_res_blocks: int = 2
+    z_channels: int = 32
+
 
 def swish(x: Tensor) -> Tensor:
     return x * torch.sigmoid(x)
@@ -341,10 +352,13 @@ class AutoEncoder(nn.Module):
             num_res_blocks=params.num_res_blocks,
             z_channels=params.z_channels,
         )
+        decoder_ch = params.ch
+        if hasattr(params, "ch_encoder"):
+            decoder_ch = params.ch_encoder
         self.decoder = Decoder(
             resolution=params.resolution,
             in_channels=params.in_channels,
-            ch=params.ch,
+            ch=decoder_ch,
             out_ch=params.out_ch,
             ch_mult=params.ch_mult,
             num_res_blocks=params.num_res_blocks,