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

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

@@ -140,7 +140,6 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
         # Some weights of the model checkpoint were not used when initializing CLIPTextModelWithProjection:
         # ['text_model.embeddings.position_ids']
 
-        self._logger.info(f"Loading model from single file at {config.path} using {load_class.__name__}")
         with SilenceWarnings():
             pipeline = load_class.from_single_file(config.path, torch_dtype=self._torch_dtype)
 

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

@@ -576,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():
@@ -589,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,
@@ -625,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,
@@ -771,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")
@@ -783,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")
@@ -800,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",
@@ -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,
@@ -898,10 +1019,17 @@ zimage_bundle: list[StarterModel] = [
     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),
 }
 

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/util/vae_working_memory.py

@@ -47,8 +47,6 @@ def estimate_vae_working_memory_sd15_sdxl(
         # If we are running in FP32, then we should account for the likely increase in model size (~250MB).
         working_memory += 250 * 2**20
 
-    print(f"estimate_vae_working_memory_sd15_sdxl: {int(working_memory)}")
-
     return int(working_memory)