InvokeAI 6.10.0rc1__py3-none-any.whl → 6.11.0__py3-none-any.whl

invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py

@@ -37,12 +37,14 @@ class GenericDiffusersLoader(ModelLoader):
         repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None
         variant = repo_variant.value if repo_variant else None
         try:
-            result: AnyModel = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, variant=variant)
+            result: AnyModel = model_class.from_pretrained(
+                model_path, torch_dtype=self._torch_dtype, variant=variant, local_files_only=True
+            )
         except OSError as e:
             if variant and "no file named" in str(
                 e
             ):  # try without the variant, just in case user's preferences changed
-                result = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype)
+                result = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, local_files_only=True)
             else:
                 raise e
         return result

invokeai/backend/model_manager/load/model_loaders/onnx.py

@@ -38,5 +38,6 @@ class OnnyxDiffusersModel(GenericDiffusersLoader):
             model_path,
             torch_dtype=self._torch_dtype,
             variant=variant,
+            local_files_only=True,
         )
         return result

invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py

@@ -80,12 +80,13 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
                 model_path,
                 torch_dtype=self._torch_dtype,
                 variant=variant,
+                local_files_only=True,
             )
         except OSError as e:
             if variant and "no file named" in str(
                 e
             ):  # try without the variant, just in case user's preferences changed
-                result = load_class.from_pretrained(model_path, torch_dtype=self._torch_dtype)
+                result = load_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, local_files_only=True)
             else:
                 raise e
 

invokeai/backend/model_manager/load/model_loaders/z_image.py

@@ -384,15 +384,19 @@ class Qwen3EncoderLoader(ModelLoader):
 
         match submodel_type:
             case SubModelType.Tokenizer:
-                return AutoTokenizer.from_pretrained(tokenizer_path)
+                # Use local_files_only=True to prevent network requests for validation
+                # The tokenizer files should already exist locally in the model directory
+                return AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
             case SubModelType.TextEncoder:
                 # Determine safe dtype based on target device capabilities
                 target_device = TorchDevice.choose_torch_device()
                 model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+                # Use local_files_only=True to prevent network requests for validation
                 return Qwen3ForCausalLM.from_pretrained(
                     text_encoder_path,
                     torch_dtype=model_dtype,
                     low_cpu_mem_usage=True,
+                    local_files_only=True,
                 )
 
         raise ValueError(
@@ -526,12 +530,27 @@ class Qwen3EncoderCheckpointLoader(ModelLoader):
                 return self._load_from_singlefile(config)
             case SubModelType.Tokenizer:
                 # For single-file Qwen3, load tokenizer from HuggingFace
-                return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE)
+                # Try local cache first to support offline usage after initial download
+                return self._load_tokenizer_with_offline_fallback()
 
         raise ValueError(
             f"Only TextEncoder and Tokenizer submodels are supported. Received: {submodel_type.value if submodel_type else 'None'}"
         )
 
+    def _load_tokenizer_with_offline_fallback(self) -> AnyModel:
+        """Load tokenizer with local_files_only fallback for offline support.
+
+        First tries to load from local cache (offline), falling back to network download
+        if the tokenizer hasn't been cached yet. This ensures offline operation after
+        the initial download.
+        """
+        try:
+            # Try loading from local cache first (supports offline usage)
+            return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE, local_files_only=True)
+        except OSError:
+            # Not in cache yet, download from HuggingFace
+            return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE)
+
     def _load_from_singlefile(
         self,
         config: AnyModelConfig,
@@ -557,7 +576,54 @@ class Qwen3EncoderCheckpointLoader(ModelLoader):
         # Load the state dict from safetensors file
         sd = load_file(model_path)
 
-        # Determine Qwen model configuration from state dict
+        # Handle ComfyUI quantized checkpoints
+        # ComfyUI stores quantized weights with accompanying scale factors:
+        # - layer.weight: quantized data (FP8)
+        # - layer.weight_scale: scale factor (FP32 scalar)
+        # Dequantization formula: dequantized = weight.to(dtype) * weight_scale
+        # Reference: https://github.com/Comfy-Org/ComfyUI/blob/master/QUANTIZATION.md
+        original_key_count = len(sd)
+        weight_scale_keys = [k for k in sd.keys() if k.endswith(".weight_scale")]
+        dequantized_count = 0
+
+        for scale_key in weight_scale_keys:
+            # Get the corresponding weight key (remove "_scale" suffix)
+            weight_key = scale_key.replace(".weight_scale", ".weight")
+            if weight_key in sd:
+                weight = sd[weight_key]
+                scale = sd[scale_key]
+                # Dequantize: convert to float and multiply by scale
+                # Handle block-wise quantization (e.g., FP4 with block_size=8)
+                # where scale has shape [weight_dim / block_size, ...]
+                # Note: Float8 types (e.g., float8_e4m3fn) require .float() instead of .to(torch.float32)
+                # as PyTorch doesn't support direct type promotion for Float8 types
+                weight_float = weight.float()
+                scale = scale.float()
+                if scale.shape != weight_float.shape and scale.numel() > 1:
+                    # Block-wise quantization: need to expand scale to match weight shape
+                    # Find which dimension differs and repeat scale along that dimension
+                    for dim in range(len(weight_float.shape)):
+                        if dim < len(scale.shape) and scale.shape[dim] != weight_float.shape[dim]:
+                            block_size = weight_float.shape[dim] // scale.shape[dim]
+                            if block_size > 1:
+                                # Repeat scale along this dimension to match weight shape
+                                scale = scale.repeat_interleave(block_size, dim=dim)
+                sd[weight_key] = weight_float * scale
+                dequantized_count += 1
+
+        if dequantized_count > 0:
+            logger.info(f"Dequantized {dequantized_count} ComfyUI quantized weights")
+
+        # Filter out ComfyUI quantization metadata keys (comfy_quant, weight_scale)
+        # These are no longer needed after dequantization
+        comfy_metadata_keys = [k for k in sd.keys() if "comfy_quant" in k or "weight_scale" in k]
+        for k in comfy_metadata_keys:
+            del sd[k]
+        if comfy_metadata_keys:
+            logger.info(f"Filtered out {len(comfy_metadata_keys)} ComfyUI quantization metadata keys")
+
+        logger.info(f"Loaded state dict with {len(sd)} keys (originally {original_key_count})")
+
         # Count the number of layers by looking at layer keys
         layer_count = 0
         for key in sd.keys():
@@ -570,34 +636,63 @@ class Qwen3EncoderCheckpointLoader(ModelLoader):
                     except ValueError:
                         pass
 
-        # Get hidden size from embed_tokens weight shape
+        # Get vocab size from embed_tokens weight shape
         embed_weight = sd.get("model.embed_tokens.weight")
         if embed_weight is None:
             raise ValueError("Could not find model.embed_tokens.weight in state dict")
-        if embed_weight.ndim != 2:
-            raise ValueError(
-                f"Expected 2D embed_tokens weight tensor, got shape {embed_weight.shape}. "
-                "The model file may be corrupted or incompatible."
-            )
-        hidden_size = embed_weight.shape[1]
+
         vocab_size = embed_weight.shape[0]
+        embed_hidden_size = embed_weight.shape[1]
+
+        # Detect model variant based on embed_tokens hidden size and layer count
+        # FLUX 2 Klein / Z-Image uses Qwen3 configurations from ComfyUI:
+        # Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/text_encoders/llama.py
+        # - Qwen3-4B: hidden_size=2560, 36 layers, 32 heads, 8 KV heads, intermediate=9728
+        # - Qwen3-8B: hidden_size=4096, 36 layers, 32 heads, 8 KV heads, intermediate=12288
+        if embed_hidden_size == 2560 and layer_count == 36:
+            # Qwen3-4B variant (FLUX 2 Klein / Z-Image)
+            logger.info("Detected Qwen3-4B variant (FLUX 2 Klein / Z-Image)")
+            hidden_size = 2560
+            num_attention_heads = 32
+            num_kv_heads = 8
+            intermediate_size = 9728
+            head_dim = 128
+            max_position_embeddings = 40960
+        elif embed_hidden_size == 4096 and layer_count == 36:
+            # Qwen3-8B variant
+            logger.info("Detected Qwen3-8B variant")
+            hidden_size = 4096
+            num_attention_heads = 32
+            num_kv_heads = 8
+            intermediate_size = 12288
+            head_dim = 128
+            max_position_embeddings = 40960
+        else:
+            # Unknown variant - try to detect from weights
+            logger.warning(
+                f"Unknown Qwen3 variant: embed_hidden_size={embed_hidden_size}, layers={layer_count}. "
+                "Attempting to detect configuration from weights..."
+            )
+            q_proj_weight = sd.get("model.layers.0.self_attn.q_proj.weight")
+            k_proj_weight = sd.get("model.layers.0.self_attn.k_proj.weight")
+            gate_proj_weight = sd.get("model.layers.0.mlp.gate_proj.weight")
 
-        # Detect attention configuration from layer 0 weights
-        q_proj_weight = sd.get("model.layers.0.self_attn.q_proj.weight")
-        k_proj_weight = sd.get("model.layers.0.self_attn.k_proj.weight")
-        gate_proj_weight = sd.get("model.layers.0.mlp.gate_proj.weight")
+            if q_proj_weight is None or k_proj_weight is None or gate_proj_weight is None:
+                raise ValueError("Could not find attention/mlp weights to determine configuration")
 
-        if q_proj_weight is None or k_proj_weight is None or gate_proj_weight is None:
-            raise ValueError("Could not find attention/mlp weights in state dict to determine configuration")
+            hidden_size = embed_hidden_size
+            head_dim = 128
+            num_attention_heads = q_proj_weight.shape[0] // head_dim
+            num_kv_heads = k_proj_weight.shape[0] // head_dim
+            intermediate_size = gate_proj_weight.shape[0]
+            max_position_embeddings = 40960
 
-        # Calculate dimensions from actual weights
-        # Qwen3 uses head_dim separately from hidden_size
-        head_dim = 128  # Standard head dimension for Qwen3 models
-        num_attention_heads = q_proj_weight.shape[0] // head_dim
-        num_kv_heads = k_proj_weight.shape[0] // head_dim
-        intermediate_size = gate_proj_weight.shape[0]
+        logger.info(
+            f"Qwen3 config: hidden_size={hidden_size}, layers={layer_count}, "
+            f"heads={num_attention_heads}, kv_heads={num_kv_heads}, intermediate={intermediate_size}"
+        )
 
-        # Create Qwen3 config - matches the diffusers text_encoder/config.json
+        # Create Qwen3 config
         qwen_config = Qwen3Config(
             vocab_size=vocab_size,
             hidden_size=hidden_size,
@@ -606,7 +701,7 @@ class Qwen3EncoderCheckpointLoader(ModelLoader):
             num_attention_heads=num_attention_heads,
             num_key_value_heads=num_kv_heads,
             head_dim=head_dim,
-            max_position_embeddings=40960,
+            max_position_embeddings=max_position_embeddings,
             rms_norm_eps=1e-6,
             tie_word_embeddings=True,
             rope_theta=1000000.0,
@@ -686,12 +781,27 @@ class Qwen3EncoderGGUFLoader(ModelLoader):
                 return self._load_from_gguf(config)
             case SubModelType.Tokenizer:
                 # For GGUF Qwen3, load tokenizer from HuggingFace
-                return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE)
+                # Try local cache first to support offline usage after initial download
+                return self._load_tokenizer_with_offline_fallback()
 
         raise ValueError(
             f"Only TextEncoder and Tokenizer submodels are supported. Received: {submodel_type.value if submodel_type else 'None'}"
         )
 
+    def _load_tokenizer_with_offline_fallback(self) -> AnyModel:
+        """Load tokenizer with local_files_only fallback for offline support.
+
+        First tries to load from local cache (offline), falling back to network download
+        if the tokenizer hasn't been cached yet. This ensures offline operation after
+        the initial download.
+        """
+        try:
+            # Try loading from local cache first (supports offline usage)
+            return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE, local_files_only=True)
+        except OSError:
+            # Not in cache yet, download from HuggingFace
+            return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE)
+
     def _load_from_gguf(
         self,
         config: AnyModelConfig,
@@ -737,7 +847,7 @@ class Qwen3EncoderGGUFLoader(ModelLoader):
                     except ValueError:
                         pass
 
-        # Get hidden size from embed_tokens weight shape
+        # Get vocab size from embed_tokens weight shape
         embed_weight = sd.get("model.embed_tokens.weight")
         if embed_weight is None:
             raise ValueError("Could not find model.embed_tokens.weight in state dict")
@@ -749,13 +859,23 @@ class Qwen3EncoderGGUFLoader(ModelLoader):
                 f"Expected 2D embed_tokens weight tensor, got shape {embed_shape}. "
                 "The model file may be corrupted or incompatible."
             )
-        hidden_size = embed_shape[1]
         vocab_size = embed_shape[0]
 
-        # Detect attention configuration from layer 0 weights
-        q_proj_weight = sd.get("model.layers.0.self_attn.q_proj.weight")
-        k_proj_weight = sd.get("model.layers.0.self_attn.k_proj.weight")
-        gate_proj_weight = sd.get("model.layers.0.mlp.gate_proj.weight")
+        # Detect attention configuration from layer weights
+        # IMPORTANT: Use layer 1 (not layer 0) because some models like FLUX 2 Klein have a special
+        # first layer with different dimensions (input projection layer) while the rest of the
+        # transformer layers have a different hidden_size. Using a middle layer ensures we get
+        # the representative hidden_size for the bulk of the model.
+        # Fall back to layer 0 if layer 1 doesn't exist.
+        q_proj_weight = sd.get("model.layers.1.self_attn.q_proj.weight")
+        k_proj_weight = sd.get("model.layers.1.self_attn.k_proj.weight")
+        gate_proj_weight = sd.get("model.layers.1.mlp.gate_proj.weight")
+
+        # Fall back to layer 0 if layer 1 doesn't exist (single-layer model edge case)
+        if q_proj_weight is None:
+            q_proj_weight = sd.get("model.layers.0.self_attn.q_proj.weight")
+            k_proj_weight = sd.get("model.layers.0.self_attn.k_proj.weight")
+            gate_proj_weight = sd.get("model.layers.0.mlp.gate_proj.weight")
 
         if q_proj_weight is None or k_proj_weight is None or gate_proj_weight is None:
             raise ValueError("Could not find attention/mlp weights in state dict to determine configuration")
@@ -766,7 +886,14 @@ class Qwen3EncoderGGUFLoader(ModelLoader):
         gate_shape = gate_proj_weight.shape if hasattr(gate_proj_weight, "shape") else gate_proj_weight.tensor_shape
 
         # Calculate dimensions from actual weights
+        # IMPORTANT: Use hidden_size from k_proj input dimension (not q_proj or embed_tokens).
+        # Some models (like FLUX 2 Klein) have unusual architectures where:
+        # - embed_tokens has a larger dimension (e.g., 2560)
+        # - q_proj may have a larger input dimension for query expansion
+        # - k_proj/v_proj have the actual transformer hidden_size (e.g., 1280)
+        # Using k_proj ensures we get the correct internal hidden_size.
         head_dim = 128  # Standard head dimension for Qwen3 models
+        hidden_size = k_shape[1]  # Use k_proj input dim as the hidden_size
         num_attention_heads = q_shape[0] // head_dim
         num_kv_heads = k_shape[0] // head_dim
         intermediate_size = gate_shape[0]

invokeai/backend/model_manager/starter_models.py

@@ -690,6 +690,115 @@ flux_fill = StarterModel(
 )
 # endregion
 
+# region FLUX.2 Klein
+flux2_vae = StarterModel(
+    name="FLUX.2 VAE",
+    base=BaseModelType.Flux2,
+    source="black-forest-labs/FLUX.2-klein-4B::vae",
+    description="FLUX.2 VAE (16-channel, same architecture as FLUX.1 VAE). ~335MB",
+    type=ModelType.VAE,
+)
+
+flux2_klein_qwen3_4b_encoder = StarterModel(
+    name="FLUX.2 Klein Qwen3 4B Encoder",
+    base=BaseModelType.Any,
+    source="black-forest-labs/FLUX.2-klein-4B::text_encoder+tokenizer",
+    description="Qwen3 4B text encoder for FLUX.2 Klein 4B (also compatible with Z-Image). ~8GB",
+    type=ModelType.Qwen3Encoder,
+)
+
+flux2_klein_qwen3_8b_encoder = StarterModel(
+    name="FLUX.2 Klein Qwen3 8B Encoder",
+    base=BaseModelType.Any,
+    source="black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizer",
+    description="Qwen3 8B text encoder for FLUX.2 Klein 9B models. ~16GB",
+    type=ModelType.Qwen3Encoder,
+)
+
+flux2_klein_4b = StarterModel(
+    name="FLUX.2 Klein 4B (Diffusers)",
+    base=BaseModelType.Flux2,
+    source="black-forest-labs/FLUX.2-klein-4B",
+    description="FLUX.2 Klein 4B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~10GB",
+    type=ModelType.Main,
+)
+
+flux2_klein_4b_single = StarterModel(
+    name="FLUX.2 Klein 4B",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/resolve/main/flux-2-klein-4b.safetensors",
+    description="FLUX.2 Klein 4B standalone transformer. Installs with VAE and Qwen3 4B encoder. ~8GB",
+    type=ModelType.Main,
+    dependencies=[flux2_vae, flux2_klein_qwen3_4b_encoder],
+)
+
+flux2_klein_4b_fp8 = StarterModel(
+    name="FLUX.2 Klein 4B (FP8)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/black-forest-labs/FLUX.2-klein-4b-fp8/resolve/main/flux-2-klein-4b-fp8.safetensors",
+    description="FLUX.2 Klein 4B FP8 quantized - smaller and faster. Installs with VAE and Qwen3 4B encoder. ~4GB",
+    type=ModelType.Main,
+    dependencies=[flux2_vae, flux2_klein_qwen3_4b_encoder],
+)
+
+flux2_klein_9b = StarterModel(
+    name="FLUX.2 Klein 9B (Diffusers)",
+    base=BaseModelType.Flux2,
+    source="black-forest-labs/FLUX.2-klein-9B",
+    description="FLUX.2 Klein 9B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~20GB",
+    type=ModelType.Main,
+)
+
+flux2_klein_9b_fp8 = StarterModel(
+    name="FLUX.2 Klein 9B (FP8)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8/resolve/main/flux-2-klein-9b-fp8.safetensors",
+    description="FLUX.2 Klein 9B FP8 quantized - more efficient than full precision. Installs with VAE and Qwen3 8B encoder. ~9.5GB",
+    type=ModelType.Main,
+    dependencies=[flux2_vae, flux2_klein_qwen3_8b_encoder],
+)
+
+flux2_klein_4b_gguf_q4 = StarterModel(
+    name="FLUX.2 Klein 4B (GGUF Q4)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/unsloth/FLUX.2-klein-4B-GGUF/resolve/main/flux-2-klein-4b-Q4_K_M.gguf",
+    description="FLUX.2 Klein 4B GGUF Q4_K_M quantized - runs on 6-8GB VRAM. Installs with VAE and Qwen3 4B encoder. ~2.6GB",
+    type=ModelType.Main,
+    format=ModelFormat.GGUFQuantized,
+    dependencies=[flux2_vae, flux2_klein_qwen3_4b_encoder],
+)
+
+flux2_klein_4b_gguf_q8 = StarterModel(
+    name="FLUX.2 Klein 4B (GGUF Q8)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/unsloth/FLUX.2-klein-4B-GGUF/resolve/main/flux-2-klein-4b-Q8_0.gguf",
+    description="FLUX.2 Klein 4B GGUF Q8_0 quantized - higher quality than Q4. Installs with VAE and Qwen3 4B encoder. ~4.3GB",
+    type=ModelType.Main,
+    format=ModelFormat.GGUFQuantized,
+    dependencies=[flux2_vae, flux2_klein_qwen3_4b_encoder],
+)
+
+flux2_klein_9b_gguf_q4 = StarterModel(
+    name="FLUX.2 Klein 9B (GGUF Q4)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/unsloth/FLUX.2-klein-9B-GGUF/resolve/main/flux-2-klein-9b-Q4_K_M.gguf",
+    description="FLUX.2 Klein 9B GGUF Q4_K_M quantized - runs on 12GB+ VRAM. Installs with VAE and Qwen3 8B encoder. ~5.8GB",
+    type=ModelType.Main,
+    format=ModelFormat.GGUFQuantized,
+    dependencies=[flux2_vae, flux2_klein_qwen3_8b_encoder],
+)
+
+flux2_klein_9b_gguf_q8 = StarterModel(
+    name="FLUX.2 Klein 9B (GGUF Q8)",
+    base=BaseModelType.Flux2,
+    source="https://huggingface.co/unsloth/FLUX.2-klein-9B-GGUF/resolve/main/flux-2-klein-9b-Q8_0.gguf",
+    description="FLUX.2 Klein 9B GGUF Q8_0 quantized - higher quality than Q4. Installs with VAE and Qwen3 8B encoder. ~10GB",
+    type=ModelType.Main,
+    format=ModelFormat.GGUFQuantized,
+    dependencies=[flux2_vae, flux2_klein_qwen3_8b_encoder],
+)
+# endregion
+
 # region Z-Image
 z_image_qwen3_encoder = StarterModel(
     name="Z-Image Qwen3 Text Encoder",
@@ -720,20 +829,20 @@ z_image_turbo_quantized = StarterModel(
     name="Z-Image Turbo (quantized)",
     base=BaseModelType.ZImage,
     source="https://huggingface.co/leejet/Z-Image-Turbo-GGUF/resolve/main/z_image_turbo-Q4_K.gguf",
-    description="Z-Image Turbo quantized to GGUF Q4_K format. Requires separate Qwen3 text encoder. ~4GB",
+    description="Z-Image Turbo quantized to GGUF Q4_K format. Requires standalone Qwen3 text encoder and Flux VAE. ~4GB",
     type=ModelType.Main,
     format=ModelFormat.GGUFQuantized,
-    dependencies=[z_image_qwen3_encoder_quantized],
+    dependencies=[z_image_qwen3_encoder_quantized, flux_vae],
 )
 
 z_image_turbo_q8 = StarterModel(
     name="Z-Image Turbo (Q8)",
     base=BaseModelType.ZImage,
     source="https://huggingface.co/leejet/Z-Image-Turbo-GGUF/resolve/main/z_image_turbo-Q8_0.gguf",
-    description="Z-Image Turbo quantized to GGUF Q8_0 format. Higher quality, larger size. Requires separate Qwen3 text encoder. ~6.6GB",
+    description="Z-Image Turbo quantized to GGUF Q8_0 format. Higher quality, larger size. Requires standalone Qwen3 text encoder and Flux VAE. ~6.6GB",
     type=ModelType.Main,
     format=ModelFormat.GGUFQuantized,
-    dependencies=[z_image_qwen3_encoder_quantized],
+    dependencies=[z_image_qwen3_encoder_quantized, flux_vae],
 )
 
 z_image_controlnet_union = StarterModel(
@@ -826,6 +935,18 @@ STARTER_MODELS: list[StarterModel] = [
     flux_redux,
     llava_onevision,
     flux_fill,
+    flux2_vae,
+    flux2_klein_4b,
+    flux2_klein_4b_single,
+    flux2_klein_4b_fp8,
+    flux2_klein_9b,
+    flux2_klein_9b_fp8,
+    flux2_klein_4b_gguf_q4,
+    flux2_klein_4b_gguf_q8,
+    flux2_klein_9b_gguf_q4,
+    flux2_klein_9b_gguf_q8,
+    flux2_klein_qwen3_4b_encoder,
+    flux2_klein_qwen3_8b_encoder,
     cogview4,
     flux_krea,
     flux_krea_quantized,
@@ -890,10 +1011,26 @@ flux_bundle: list[StarterModel] = [
     flux_krea_quantized,
 ]
 
+zimage_bundle: list[StarterModel] = [
+    z_image_turbo_quantized,
+    z_image_qwen3_encoder_quantized,
+    z_image_controlnet_union,
+    z_image_controlnet_tile,
+    flux_vae,
+]
+
+flux2_klein_bundle: list[StarterModel] = [
+    flux2_klein_4b_gguf_q4,
+    flux2_vae,
+    flux2_klein_qwen3_4b_encoder,
+]
+
 STARTER_BUNDLES: dict[str, StarterModelBundle] = {
     BaseModelType.StableDiffusion1: StarterModelBundle(name="Stable Diffusion 1.5", models=sd1_bundle),
     BaseModelType.StableDiffusionXL: StarterModelBundle(name="SDXL", models=sdxl_bundle),
     BaseModelType.Flux: StarterModelBundle(name="FLUX.1 dev", models=flux_bundle),
+    BaseModelType.Flux2: StarterModelBundle(name="FLUX.2 Klein", models=flux2_klein_bundle),
+    BaseModelType.ZImage: StarterModelBundle(name="Z-Image Turbo", models=zimage_bundle),
 }
 
 assert len(STARTER_MODELS) == len({m.source for m in STARTER_MODELS}), "Duplicate starter models"

invokeai/backend/model_manager/taxonomy.py

@@ -46,6 +46,8 @@ class BaseModelType(str, Enum):
     """Indicates the model is associated with the Stable Diffusion XL Refiner model architecture."""
     Flux = "flux"
     """Indicates the model is associated with FLUX.1 model architecture, including FLUX Dev, Schnell and Fill."""
+    Flux2 = "flux2"
+    """Indicates the model is associated with FLUX.2 model architecture, including FLUX2 Klein."""
     CogView4 = "cogview4"
     """Indicates the model is associated with CogView 4 model architecture."""
     ZImage = "z-image"
@@ -111,11 +113,36 @@ class ModelVariantType(str, Enum):
 
 
 class FluxVariantType(str, Enum):
+    """FLUX.1 model variants."""
+
     Schnell = "schnell"
     Dev = "dev"
     DevFill = "dev_fill"
 
 
+class Flux2VariantType(str, Enum):
+    """FLUX.2 model variants."""
+
+    Klein4B = "klein_4b"
+    """Flux2 Klein 4B variant using Qwen3 4B text encoder."""
+
+    Klein9B = "klein_9b"
+    """Flux2 Klein 9B variant using Qwen3 8B text encoder (distilled)."""
+
+    Klein9BBase = "klein_9b_base"
+    """Flux2 Klein 9B Base variant - undistilled foundation model using Qwen3 8B text encoder."""
+
+
+class Qwen3VariantType(str, Enum):
+    """Qwen3 text encoder variants based on model size."""
+
+    Qwen3_4B = "qwen3_4b"
+    """Qwen3 4B text encoder (hidden_size=2560). Used by FLUX.2 Klein 4B and Z-Image."""
+
+    Qwen3_8B = "qwen3_8b"
+    """Qwen3 8B text encoder (hidden_size=4096). Used by FLUX.2 Klein 9B."""
+
+
 class ModelFormat(str, Enum):
     """Storage format of model."""
 
@@ -174,7 +201,7 @@ class FluxLoRAFormat(str, Enum):
     XLabs = "flux.xlabs"
 
 
-AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, FluxVariantType]
-variant_type_adapter = TypeAdapter[ModelVariantType | ClipVariantType | FluxVariantType](
-    ModelVariantType | ClipVariantType | FluxVariantType
-)
+AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, Qwen3VariantType]
+variant_type_adapter = TypeAdapter[
+    ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType
+](ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType)

invokeai/backend/model_manager/util/select_hf_files.py

@@ -60,7 +60,7 @@ def filter_files(
 
     # Start by filtering on model file extensions, discarding images, docs, etc
     for file in files:
-        if file.name.endswith((".json", ".txt")):
+        if file.name.endswith((".json", ".txt", ".jinja")):  # .jinja for chat templates
             paths.append(file)
         elif file.name.endswith(
             (
@@ -116,7 +116,8 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
 
         # Note: '.model' was added to support:
         # https://huggingface.co/black-forest-labs/FLUX.1-schnell/blob/768d12a373ed5cc9ef9a9dea7504dc09fcc14842/tokenizer_2/spiece.model
-        elif path.suffix in [".json", ".txt", ".model"]:
+        # Note: '.jinja' was added to support chat templates for FLUX.2 Klein models
+        elif path.suffix in [".json", ".txt", ".model", ".jinja"]:
             result.add(path)
 
         elif variant in [

invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py

@@ -140,16 +140,50 @@ def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | N
 
 
 def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
-    """Groups the keys in the state dict by layer."""
+    """Groups the keys in the state dict by layer.
+
+    Z-Image LoRAs have keys like:
+    - diffusion_model.layers.17.attention.to_k.alpha
+    - diffusion_model.layers.17.attention.to_k.dora_scale
+    - diffusion_model.layers.17.attention.to_k.lora_down.weight
+    - diffusion_model.layers.17.attention.to_k.lora_up.weight
+
+    We need to group these by the full layer path (e.g., diffusion_model.layers.17.attention.to_k)
+    and extract the suffix (alpha, dora_scale, lora_down.weight, lora_up.weight).
+    """
     layer_dict: dict[str, dict[str, torch.Tensor]] = {}
+
+    # Known suffixes that indicate the end of a layer name
+    known_suffixes = [
+        ".lora_A.weight",
+        ".lora_B.weight",
+        ".lora_down.weight",
+        ".lora_up.weight",
+        ".dora_scale",
+        ".alpha",
+    ]
+
     for key in state_dict:
         if not isinstance(key, str):
             continue
-        # Split the 'lora_A.weight' or 'lora_B.weight' suffix from the layer name.
-        parts = key.rsplit(".", maxsplit=2)
-        layer_name = parts[0]
-        key_name = ".".join(parts[1:])
+
+        # Try to find a known suffix
+        layer_name = None
+        key_name = None
+        for suffix in known_suffixes:
+            if key.endswith(suffix):
+                layer_name = key[: -len(suffix)]
+                key_name = suffix[1:]  # Remove leading dot
+                break
+
+        if layer_name is None:
+            # Fallback to original logic for unknown formats
+            parts = key.rsplit(".", maxsplit=2)
+            layer_name = parts[0]
+            key_name = ".".join(parts[1:])
+
         if layer_name not in layer_dict:
             layer_dict[layer_name] = {}
         layer_dict[layer_name][key_name] = state_dict[key]
+
     return layer_dict