diff --git a/invokeai/backend/model_manager/load/model_loaders/anima.py b/invokeai/backend/model_manager/load/model_loaders/anima.py index 8d068f5468c..8cdb99ad8f2 100644 --- a/invokeai/backend/model_manager/load/model_loaders/anima.py +++ b/invokeai/backend/model_manager/load/model_loaders/anima.py @@ -24,6 +24,35 @@ logger = InvokeAILogger.get_logger(__name__) +def _strip_anima_bundle_prefix(sd: dict) -> dict: + """Strip the transformer-key prefix from an Anima single-file checkpoint. + + Handles both packaging formats: + - Official format: keys prefixed with `net.` (e.g. `net.blocks.0...`) + - ComfyUI bundled format: transformer keys prefixed with `model.diffusion_model.` + alongside `first_stage_model.*` (VAE) and `cond_stage_model.*` (text encoder). + + Only keys under the detected prefix are kept; unrelated keys from bundled + checkpoints (VAE, text encoder) are dropped. If no known prefix is present, the + state dict is returned unchanged. + """ + prefix_to_strip = None + for prefix in ["model.diffusion_model.", "net."]: + if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)): + prefix_to_strip = prefix + break + + if prefix_to_strip is None: + return sd + + stripped_sd: dict = {} + for key, value in sd.items(): + if isinstance(key, str) and key.startswith(prefix_to_strip): + stripped_sd[key[len(prefix_to_strip) :]] = value + # Skip non-transformer keys from bundled checkpoints (VAE, text encoder) + return stripped_sd + + @ModelLoaderRegistry.register(base=BaseModelType.Anima, type=ModelType.Main, format=ModelFormat.Checkpoint) class AnimaCheckpointModel(ModelLoader): """Class to load Anima transformer models from single-file checkpoints. @@ -67,23 +96,8 @@ def _load_from_singlefile( # Load the state dict from safetensors sd = load_file(model_path) - # Handle different checkpoint packaging formats: - # - Official format: keys prefixed with `net.` (e.g. `net.blocks.0...`) - # - ComfyUI bundled format: transformer keys prefixed with `model.diffusion_model.` - # alongside `first_stage_model.*` (VAE) and `cond_stage_model.*` (text encoder) - prefix_to_strip = None - for prefix in ["model.diffusion_model.", "net."]: - if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)): - prefix_to_strip = prefix - break - - if prefix_to_strip: - stripped_sd = {} - for key, value in sd.items(): - if isinstance(key, str) and key.startswith(prefix_to_strip): - stripped_sd[key[len(prefix_to_strip) :]] = value - # Skip non-transformer keys from bundled checkpoints (VAE, text encoder) - sd = stripped_sd + # Strip the transformer-key prefix (`net.` or bundled `model.diffusion_model.`). + sd = _strip_anima_bundle_prefix(sd) # Create an empty AnimaTransformer with Anima's default architecture parameters with accelerate.init_empty_weights(): diff --git a/invokeai/backend/model_manager/load/model_loaders/flux.py b/invokeai/backend/model_manager/load/model_loaders/flux.py index c15f9f26abf..16830b9e02c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux.py @@ -54,6 +54,10 @@ from invokeai.backend.model_manager.configs.vae import VAE_Checkpoint_Config_Base, VAE_Checkpoint_Flux2_Config from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.flux2_state_dict_utils import ( + convert_flux2_bfl_to_diffusers, + convert_flux2_vae_bfl_to_diffusers, +) from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import ( AnyModel, @@ -177,7 +181,7 @@ def _load_model( for k in sd.keys() ) if is_bfl_format: - sd = self._convert_flux2_vae_bfl_to_diffusers(sd) + sd = convert_flux2_vae_bfl_to_diffusers(sd) # FLUX.2 VAE configuration (32 latent channels). # The standard FLUX.2 VAE uses block_out_channels=(128,256,512,512) for both @@ -236,171 +240,6 @@ def _load_model( model = self._apply_fp8_layerwise_casting(model, config, submodel_type) return model - def _convert_flux2_vae_bfl_to_diffusers(self, sd: dict) -> dict: - """Convert FLUX.2 VAE BFL format state dict to diffusers format. - - Key differences: - - encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y - - encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv - - encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1 - - encoder.mid.attn_1.q/k/v -> encoder.mid_block.attentions.0.to_q/k/v - - encoder.norm_out -> encoder.conv_norm_out - - encoder.quant_conv -> quant_conv (top-level) - - decoder.up.X -> decoder.up_blocks.(num_blocks-1-X) (reversed order!) - - decoder.post_quant_conv -> post_quant_conv (top-level) - - *.nin_shortcut -> *.conv_shortcut - """ - import re - - converted = {} - num_up_blocks = 4 # Standard VAE has 4 up blocks - - for old_key, tensor in sd.items(): - new_key = old_key - - # Encoder down blocks: encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y - match = re.match(r"encoder\.down\.(\d+)\.block\.(\d+)\.(.*)", old_key) - if match: - block_idx, resnet_idx, rest = match.groups() - rest = rest.replace("nin_shortcut", "conv_shortcut") - new_key = f"encoder.down_blocks.{block_idx}.resnets.{resnet_idx}.{rest}" - converted[new_key] = tensor - continue - - # Encoder downsamplers: encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv - match = re.match(r"encoder\.down\.(\d+)\.downsample\.conv\.(.*)", old_key) - if match: - block_idx, rest = match.groups() - new_key = f"encoder.down_blocks.{block_idx}.downsamplers.0.conv.{rest}" - converted[new_key] = tensor - continue - - # Encoder mid block resnets: encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1 - match = re.match(r"encoder\.mid\.block_(\d+)\.(.*)", old_key) - if match: - block_num, rest = match.groups() - resnet_idx = int(block_num) - 1 # block_1 -> resnets.0, block_2 -> resnets.1 - new_key = f"encoder.mid_block.resnets.{resnet_idx}.{rest}" - converted[new_key] = tensor - continue - - # Encoder mid block attention: encoder.mid.attn_1.* -> encoder.mid_block.attentions.0.* - match = re.match(r"encoder\.mid\.attn_1\.(.*)", old_key) - if match: - rest = match.group(1) - # Map attention keys - # BFL uses Conv2d (shape [out, in, 1, 1]), diffusers uses Linear (shape [out, in]) - # Squeeze the extra dimensions for weight tensors - if rest.startswith("q."): - new_key = f"encoder.mid_block.attentions.0.to_q.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("k."): - new_key = f"encoder.mid_block.attentions.0.to_k.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("v."): - new_key = f"encoder.mid_block.attentions.0.to_v.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("proj_out."): - new_key = f"encoder.mid_block.attentions.0.to_out.0.{rest[9:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("norm."): - new_key = f"encoder.mid_block.attentions.0.group_norm.{rest[5:]}" - else: - new_key = f"encoder.mid_block.attentions.0.{rest}" - converted[new_key] = tensor - continue - - # Encoder norm_out -> conv_norm_out - if old_key.startswith("encoder.norm_out."): - new_key = old_key.replace("encoder.norm_out.", "encoder.conv_norm_out.") - converted[new_key] = tensor - continue - - # Encoder quant_conv -> quant_conv (move to top level) - if old_key.startswith("encoder.quant_conv."): - new_key = old_key.replace("encoder.quant_conv.", "quant_conv.") - converted[new_key] = tensor - continue - - # Decoder up blocks (reversed order!): decoder.up.X -> decoder.up_blocks.(num_blocks-1-X) - match = re.match(r"decoder\.up\.(\d+)\.block\.(\d+)\.(.*)", old_key) - if match: - block_idx, resnet_idx, rest = match.groups() - # Reverse the block index - new_block_idx = num_up_blocks - 1 - int(block_idx) - rest = rest.replace("nin_shortcut", "conv_shortcut") - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.{rest}" - converted[new_key] = tensor - continue - - # Decoder upsamplers (reversed order!) - match = re.match(r"decoder\.up\.(\d+)\.upsample\.conv\.(.*)", old_key) - if match: - block_idx, rest = match.groups() - new_block_idx = num_up_blocks - 1 - int(block_idx) - new_key = f"decoder.up_blocks.{new_block_idx}.upsamplers.0.conv.{rest}" - converted[new_key] = tensor - continue - - # Decoder mid block resnets: decoder.mid.block_1/2 -> decoder.mid_block.resnets.0/1 - match = re.match(r"decoder\.mid\.block_(\d+)\.(.*)", old_key) - if match: - block_num, rest = match.groups() - resnet_idx = int(block_num) - 1 - new_key = f"decoder.mid_block.resnets.{resnet_idx}.{rest}" - converted[new_key] = tensor - continue - - # Decoder mid block attention: decoder.mid.attn_1.* -> decoder.mid_block.attentions.0.* - match = re.match(r"decoder\.mid\.attn_1\.(.*)", old_key) - if match: - rest = match.group(1) - # BFL uses Conv2d (shape [out, in, 1, 1]), diffusers uses Linear (shape [out, in]) - # Squeeze the extra dimensions for weight tensors - if rest.startswith("q."): - new_key = f"decoder.mid_block.attentions.0.to_q.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("k."): - new_key = f"decoder.mid_block.attentions.0.to_k.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("v."): - new_key = f"decoder.mid_block.attentions.0.to_v.{rest[2:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("proj_out."): - new_key = f"decoder.mid_block.attentions.0.to_out.0.{rest[9:]}" - if rest.endswith(".weight") and tensor.dim() == 4: - tensor = tensor.squeeze(-1).squeeze(-1) - elif rest.startswith("norm."): - new_key = f"decoder.mid_block.attentions.0.group_norm.{rest[5:]}" - else: - new_key = f"decoder.mid_block.attentions.0.{rest}" - converted[new_key] = tensor - continue - - # Decoder norm_out -> conv_norm_out - if old_key.startswith("decoder.norm_out."): - new_key = old_key.replace("decoder.norm_out.", "decoder.conv_norm_out.") - converted[new_key] = tensor - continue - - # Decoder post_quant_conv -> post_quant_conv (move to top level) - if old_key.startswith("decoder.post_quant_conv."): - new_key = old_key.replace("decoder.post_quant_conv.", "post_quant_conv.") - converted[new_key] = tensor - continue - - # Keep other keys as-is (like encoder.conv_in, decoder.conv_in, decoder.conv_out, bn.*) - converted[new_key] = tensor - - return converted - @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPEmbed, format=ModelFormat.Diffusers) class CLIPDiffusersLoader(ModelLoader): @@ -814,7 +653,7 @@ def _load_from_singlefile( } # Convert BFL format state dict to diffusers format - converted_sd = self._convert_flux2_bfl_to_diffusers(sd) + converted_sd = convert_flux2_bfl_to_diffusers(sd) # Detect architecture from checkpoint keys double_block_indices = [ @@ -899,156 +738,6 @@ def _load_from_singlefile( return model - def _convert_flux2_bfl_to_diffusers(self, sd: dict) -> dict: - """Convert FLUX.2 BFL format state dict to diffusers format. - - Based on diffusers convert_flux2_to_diffusers.py key mappings. - """ - converted = {} - - # Basic key renames - key_renames = { - "img_in.weight": "x_embedder.weight", - "txt_in.weight": "context_embedder.weight", - "time_in.in_layer.weight": "time_guidance_embed.timestep_embedder.linear_1.weight", - "time_in.out_layer.weight": "time_guidance_embed.timestep_embedder.linear_2.weight", - "guidance_in.in_layer.weight": "time_guidance_embed.guidance_embedder.linear_1.weight", - "guidance_in.out_layer.weight": "time_guidance_embed.guidance_embedder.linear_2.weight", - "double_stream_modulation_img.lin.weight": "double_stream_modulation_img.linear.weight", - "double_stream_modulation_txt.lin.weight": "double_stream_modulation_txt.linear.weight", - "single_stream_modulation.lin.weight": "single_stream_modulation.linear.weight", - "final_layer.linear.weight": "proj_out.weight", - "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight", - } - - for old_key, tensor in sd.items(): - new_key = old_key - - # Apply basic renames - if old_key in key_renames: - new_key = key_renames[old_key] - # Apply scale-shift swap for adaLN modulation weights - # BFL and diffusers use different parameter ordering for AdaLayerNorm - if old_key == "final_layer.adaLN_modulation.1.weight": - tensor = self._swap_scale_shift(tensor) - converted[new_key] = tensor - continue - - # Convert double_blocks.X.* to transformer_blocks.X.* - if old_key.startswith("double_blocks."): - new_key = self._convert_double_block_key(old_key, tensor, converted) - if new_key is None: - continue # Key was handled specially - # Convert single_blocks.X.* to single_transformer_blocks.X.* - elif old_key.startswith("single_blocks."): - new_key = self._convert_single_block_key(old_key, tensor, converted) - if new_key is None: - continue # Key was handled specially - - if new_key != old_key or new_key not in converted: - converted[new_key] = tensor - - return converted - - def _convert_double_block_key(self, key: str, tensor: torch.Tensor, converted: dict) -> str | None: - """Convert double_blocks key to transformer_blocks format.""" - parts = key.split(".") - block_idx = parts[1] - rest = ".".join(parts[2:]) - - prefix = f"transformer_blocks.{block_idx}" - - # Attention QKV conversion - BFL uses fused qkv, diffusers uses separate - if "img_attn.qkv.weight" in rest: - # Split fused QKV into separate Q, K, V - # Defensive check: ensure tensor has at least 1 dimension and can be split into 3 - if tensor.dim() < 1 or tensor.shape[0] % 3 != 0: - # Skip malformed tensors (might be metadata or corrupted) - return key - q, k, v = tensor.chunk(3, dim=0) - converted[f"{prefix}.attn.to_q.weight"] = q - converted[f"{prefix}.attn.to_k.weight"] = k - converted[f"{prefix}.attn.to_v.weight"] = v - return None - elif "txt_attn.qkv.weight" in rest: - # Defensive check - if tensor.dim() < 1 or tensor.shape[0] % 3 != 0: - return key - q, k, v = tensor.chunk(3, dim=0) - converted[f"{prefix}.attn.add_q_proj.weight"] = q - converted[f"{prefix}.attn.add_k_proj.weight"] = k - converted[f"{prefix}.attn.add_v_proj.weight"] = v - return None - - # Attention output projection - if "img_attn.proj.weight" in rest: - return f"{prefix}.attn.to_out.0.weight" - elif "txt_attn.proj.weight" in rest: - return f"{prefix}.attn.to_add_out.weight" - - # Attention norms - if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_q.weight" - elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_k.weight" - elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_added_q.weight" - elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_added_k.weight" - - # MLP layers - if "img_mlp.0.weight" in rest: - return f"{prefix}.ff.linear_in.weight" - elif "img_mlp.2.weight" in rest: - return f"{prefix}.ff.linear_out.weight" - elif "txt_mlp.0.weight" in rest: - return f"{prefix}.ff_context.linear_in.weight" - elif "txt_mlp.2.weight" in rest: - return f"{prefix}.ff_context.linear_out.weight" - - return key - - def _convert_single_block_key(self, key: str, tensor: torch.Tensor, converted: dict) -> str | None: - """Convert single_blocks key to single_transformer_blocks format.""" - parts = key.split(".") - block_idx = parts[1] - rest = ".".join(parts[2:]) - - prefix = f"single_transformer_blocks.{block_idx}" - - # linear1 is the fused QKV+MLP projection - if "linear1.weight" in rest: - return f"{prefix}.attn.to_qkv_mlp_proj.weight" - elif "linear2.weight" in rest: - return f"{prefix}.attn.to_out.weight" - - # Norms - if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_q.weight" - elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_k.weight" - - return key - - def _swap_scale_shift(self, weight: torch.Tensor) -> torch.Tensor: - """Swap scale and shift in AdaLayerNorm weights. - - BFL and diffusers use different parameter ordering for AdaLayerNorm. - This function swaps the two halves of the weight tensor. - - Args: - weight: Weight tensor of shape (out_features,) or (out_features, in_features) - - Returns: - Weight tensor with scale and shift swapped. - """ - # Defensive check: ensure tensor can be split - if weight.dim() < 1 or weight.shape[0] % 2 != 0: - return weight - # Split in half along the first dimension and swap - shift, scale = weight.chunk(2, dim=0) - return torch.cat([scale, shift], dim=0) - def _dequantize_fp8_weights(self, sd: dict) -> dict: """Dequantize FP8 quantized weights in the state dict. @@ -1164,7 +853,7 @@ def _load_from_singlefile( } # Convert BFL format state dict to diffusers format - converted_sd = self._convert_flux2_bfl_to_diffusers(sd) + converted_sd = convert_flux2_bfl_to_diffusers(sd) # Detect architecture from checkpoint keys double_block_indices = [ @@ -1265,127 +954,6 @@ def _load_from_singlefile( model.load_state_dict(converted_sd, assign=True) return model - def _convert_flux2_bfl_to_diffusers(self, sd: dict) -> dict: - """Convert FLUX.2 BFL format state dict to diffusers format.""" - converted = {} - - key_renames = { - "img_in.weight": "x_embedder.weight", - "txt_in.weight": "context_embedder.weight", - "time_in.in_layer.weight": "time_guidance_embed.timestep_embedder.linear_1.weight", - "time_in.out_layer.weight": "time_guidance_embed.timestep_embedder.linear_2.weight", - "guidance_in.in_layer.weight": "time_guidance_embed.guidance_embedder.linear_1.weight", - "guidance_in.out_layer.weight": "time_guidance_embed.guidance_embedder.linear_2.weight", - "double_stream_modulation_img.lin.weight": "double_stream_modulation_img.linear.weight", - "double_stream_modulation_txt.lin.weight": "double_stream_modulation_txt.linear.weight", - "single_stream_modulation.lin.weight": "single_stream_modulation.linear.weight", - "final_layer.linear.weight": "proj_out.weight", - "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight", - } - - for old_key, tensor in sd.items(): - new_key = old_key - - if old_key in key_renames: - new_key = key_renames[old_key] - if old_key == "final_layer.adaLN_modulation.1.weight": - tensor = self._swap_scale_shift(tensor) - converted[new_key] = tensor - continue - - if old_key.startswith("double_blocks."): - new_key = self._convert_double_block_key(old_key, tensor, converted) - if new_key is None: - continue - elif old_key.startswith("single_blocks."): - new_key = self._convert_single_block_key(old_key, tensor, converted) - if new_key is None: - continue - - if new_key != old_key or new_key not in converted: - converted[new_key] = tensor - - return converted - - def _convert_double_block_key(self, key: str, tensor, converted: dict) -> str | None: - parts = key.split(".") - block_idx = parts[1] - rest = ".".join(parts[2:]) - prefix = f"transformer_blocks.{block_idx}" - - if "img_attn.qkv.weight" in rest: - q, k, v = self._chunk_tensor(tensor, 3) - converted[f"{prefix}.attn.to_q.weight"] = q - converted[f"{prefix}.attn.to_k.weight"] = k - converted[f"{prefix}.attn.to_v.weight"] = v - return None - elif "txt_attn.qkv.weight" in rest: - q, k, v = self._chunk_tensor(tensor, 3) - converted[f"{prefix}.attn.add_q_proj.weight"] = q - converted[f"{prefix}.attn.add_k_proj.weight"] = k - converted[f"{prefix}.attn.add_v_proj.weight"] = v - return None - - if "img_attn.proj.weight" in rest: - return f"{prefix}.attn.to_out.0.weight" - elif "txt_attn.proj.weight" in rest: - return f"{prefix}.attn.to_add_out.weight" - - if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_q.weight" - elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_k.weight" - elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_added_q.weight" - elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_added_k.weight" - - if "img_mlp.0.weight" in rest: - return f"{prefix}.ff.linear_in.weight" - elif "img_mlp.2.weight" in rest: - return f"{prefix}.ff.linear_out.weight" - elif "txt_mlp.0.weight" in rest: - return f"{prefix}.ff_context.linear_in.weight" - elif "txt_mlp.2.weight" in rest: - return f"{prefix}.ff_context.linear_out.weight" - - return key - - def _convert_single_block_key(self, key: str, tensor, converted: dict) -> str | None: - parts = key.split(".") - block_idx = parts[1] - rest = ".".join(parts[2:]) - prefix = f"single_transformer_blocks.{block_idx}" - - if "linear1.weight" in rest: - return f"{prefix}.attn.to_qkv_mlp_proj.weight" - elif "linear2.weight" in rest: - return f"{prefix}.attn.to_out.weight" - - if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest: - return f"{prefix}.attn.norm_q.weight" - elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest: - return f"{prefix}.attn.norm_k.weight" - - return key - - def _chunk_tensor(self, tensor, chunks: int): - """Chunk a tensor, handling both regular tensors and GGUF quantized tensors.""" - if hasattr(tensor, "get_dequantized_tensor"): - # GGUF quantized tensor - dequantize first, then chunk - # This loses quantization for the split weights, but is necessary - # because diffusers uses separate Q/K/V projections - tensor = tensor.get_dequantized_tensor() - return tensor.chunk(chunks, dim=0) - - def _swap_scale_shift(self, weight) -> torch.Tensor: - """Swap scale and shift in AdaLayerNorm weights.""" - if hasattr(weight, "get_dequantized_tensor"): - # For GGUF, dequantize first - weight = weight.get_dequantized_tensor() - shift, scale = weight.chunk(2, dim=0) - return torch.cat([scale, shift], dim=0) - @ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.ControlNet, format=ModelFormat.Checkpoint) @ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.ControlNet, format=ModelFormat.Diffusers) diff --git a/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py b/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py new file mode 100644 index 00000000000..86fd8478169 --- /dev/null +++ b/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py @@ -0,0 +1,326 @@ +"""Pure state-dict conversion helpers for FLUX.2 single-file checkpoints. + +These functions translate the BFL (black-forest-labs) key layout used by FLUX.2 +single-file transformer and VAE checkpoints into the diffusers layout that the +`Flux2Transformer2DModel` / `AutoencoderKLFlux2` architectures expect. + +They are intentionally free of any file/model-loading side effects so the key +remapping can be unit-tested against a synthetic state dict (see +`tests/backend/model_manager/load/`). Both the checkpoint and GGUF FLUX.2 loaders +delegate here; GGUF quantized tensors are dequantized only where a fused weight +must be split (diffusers uses separate Q/K/V projections). + +Based on the diffusers `convert_flux2_to_diffusers.py` key mappings. +""" + +import re + +import torch + + +def _flux2_chunk_tensor(tensor, chunks: int): + """Chunk a tensor along dim 0, dequantizing GGUF tensors first. + + diffusers uses separate Q/K/V projections, so a fused GGUF weight cannot stay + quantized through the split. + """ + if hasattr(tensor, "get_dequantized_tensor"): + tensor = tensor.get_dequantized_tensor() + return tensor.chunk(chunks, dim=0) + + +def _flux2_malformed_for_chunk(tensor, chunks: int) -> bool: + """Return True if a plain tensor cannot be evenly split into `chunks` along dim 0. + + GGUF quantized tensors are always considered well-formed here (their logical + shape is only known after dequantization, matching the original GGUF loader, + which split them unconditionally). + """ + if hasattr(tensor, "get_dequantized_tensor"): + return False + return tensor.dim() < 1 or tensor.shape[0] % chunks != 0 + + +def _flux2_swap_scale_shift(weight): + """Swap the scale/shift halves of an AdaLayerNorm weight (BFL vs diffusers order).""" + if hasattr(weight, "get_dequantized_tensor"): + weight = weight.get_dequantized_tensor() + elif weight.dim() < 1 or weight.shape[0] % 2 != 0: + # Defensive: leave malformed plain tensors untouched. + return weight + shift, scale = weight.chunk(2, dim=0) + return torch.cat([scale, shift], dim=0) + + +def _convert_flux2_double_block_key(key: str, tensor, converted: dict) -> str | None: + """Convert a `double_blocks.X.*` key to `transformer_blocks.X.*` format. + + Returns the new key, or None if the key was consumed by writing directly into + `converted` (fused QKV split into separate projections). + """ + parts = key.split(".") + block_idx = parts[1] + rest = ".".join(parts[2:]) + + prefix = f"transformer_blocks.{block_idx}" + + # Attention QKV: BFL uses fused qkv, diffusers uses separate Q/K/V. + if "img_attn.qkv.weight" in rest: + if _flux2_malformed_for_chunk(tensor, 3): + return key + q, k, v = _flux2_chunk_tensor(tensor, 3) + converted[f"{prefix}.attn.to_q.weight"] = q + converted[f"{prefix}.attn.to_k.weight"] = k + converted[f"{prefix}.attn.to_v.weight"] = v + return None + elif "txt_attn.qkv.weight" in rest: + if _flux2_malformed_for_chunk(tensor, 3): + return key + q, k, v = _flux2_chunk_tensor(tensor, 3) + converted[f"{prefix}.attn.add_q_proj.weight"] = q + converted[f"{prefix}.attn.add_k_proj.weight"] = k + converted[f"{prefix}.attn.add_v_proj.weight"] = v + return None + + # Attention output projection + if "img_attn.proj.weight" in rest: + return f"{prefix}.attn.to_out.0.weight" + elif "txt_attn.proj.weight" in rest: + return f"{prefix}.attn.to_add_out.weight" + + # Attention norms + if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest: + return f"{prefix}.attn.norm_q.weight" + elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest: + return f"{prefix}.attn.norm_k.weight" + elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest: + return f"{prefix}.attn.norm_added_q.weight" + elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest: + return f"{prefix}.attn.norm_added_k.weight" + + # MLP layers + if "img_mlp.0.weight" in rest: + return f"{prefix}.ff.linear_in.weight" + elif "img_mlp.2.weight" in rest: + return f"{prefix}.ff.linear_out.weight" + elif "txt_mlp.0.weight" in rest: + return f"{prefix}.ff_context.linear_in.weight" + elif "txt_mlp.2.weight" in rest: + return f"{prefix}.ff_context.linear_out.weight" + + return key + + +def _convert_flux2_single_block_key(key: str, tensor, converted: dict) -> str | None: + """Convert a `single_blocks.X.*` key to `single_transformer_blocks.X.*` format.""" + parts = key.split(".") + block_idx = parts[1] + rest = ".".join(parts[2:]) + + prefix = f"single_transformer_blocks.{block_idx}" + + # linear1 is the fused QKV+MLP projection + if "linear1.weight" in rest: + return f"{prefix}.attn.to_qkv_mlp_proj.weight" + elif "linear2.weight" in rest: + return f"{prefix}.attn.to_out.weight" + + # Norms + if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest: + return f"{prefix}.attn.norm_q.weight" + elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest: + return f"{prefix}.attn.norm_k.weight" + + return key + + +def convert_flux2_bfl_to_diffusers(sd: dict) -> dict: + """Convert a FLUX.2 transformer BFL-format state dict to diffusers format.""" + converted: dict = {} + + # Basic key renames + key_renames = { + "img_in.weight": "x_embedder.weight", + "txt_in.weight": "context_embedder.weight", + "time_in.in_layer.weight": "time_guidance_embed.timestep_embedder.linear_1.weight", + "time_in.out_layer.weight": "time_guidance_embed.timestep_embedder.linear_2.weight", + "guidance_in.in_layer.weight": "time_guidance_embed.guidance_embedder.linear_1.weight", + "guidance_in.out_layer.weight": "time_guidance_embed.guidance_embedder.linear_2.weight", + "double_stream_modulation_img.lin.weight": "double_stream_modulation_img.linear.weight", + "double_stream_modulation_txt.lin.weight": "double_stream_modulation_txt.linear.weight", + "single_stream_modulation.lin.weight": "single_stream_modulation.linear.weight", + "final_layer.linear.weight": "proj_out.weight", + "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight", + } + + for old_key, tensor in sd.items(): + new_key = old_key + + # Apply basic renames + if old_key in key_renames: + new_key = key_renames[old_key] + # Apply scale-shift swap for adaLN modulation weights (BFL and diffusers use + # different parameter ordering for AdaLayerNorm). + if old_key == "final_layer.adaLN_modulation.1.weight": + tensor = _flux2_swap_scale_shift(tensor) + converted[new_key] = tensor + continue + + # Convert double_blocks.X.* to transformer_blocks.X.* + if old_key.startswith("double_blocks."): + new_key = _convert_flux2_double_block_key(old_key, tensor, converted) + if new_key is None: + continue # Key was handled specially + # Convert single_blocks.X.* to single_transformer_blocks.X.* + elif old_key.startswith("single_blocks."): + new_key = _convert_flux2_single_block_key(old_key, tensor, converted) + if new_key is None: + continue # Key was handled specially + + if new_key != old_key or new_key not in converted: + converted[new_key] = tensor + + return converted + + +def _convert_flux2_vae_mid_attention_key(rest: str, tensor, block: str): + """Map a `{enc,dec}.mid.attn_1.*` key to the diffusers mid_block attention layout. + + BFL uses Conv2d (shape `[out, in, 1, 1]`), diffusers uses Linear (`[out, in]`), so + weight tensors are squeezed. Returns `(new_key, tensor)`. + """ + attn_prefix = f"{block}.mid_block.attentions.0" + if rest.startswith("q."): + new_key = f"{attn_prefix}.to_q.{rest[2:]}" + elif rest.startswith("k."): + new_key = f"{attn_prefix}.to_k.{rest[2:]}" + elif rest.startswith("v."): + new_key = f"{attn_prefix}.to_v.{rest[2:]}" + elif rest.startswith("proj_out."): + new_key = f"{attn_prefix}.to_out.0.{rest[9:]}" + elif rest.startswith("norm."): + return f"{attn_prefix}.group_norm.{rest[5:]}", tensor + else: + return f"{attn_prefix}.{rest}", tensor + + if rest.endswith(".weight") and hasattr(tensor, "dim") and tensor.dim() == 4: + tensor = tensor.squeeze(-1).squeeze(-1) + return new_key, tensor + + +def convert_flux2_vae_bfl_to_diffusers(sd: dict) -> dict: + """Convert a FLUX.2 VAE BFL-format state dict to diffusers format. + + Key differences: + - encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y + - encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv + - encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1 + - encoder.mid.attn_1.q/k/v -> encoder.mid_block.attentions.0.to_q/k/v + - encoder.norm_out -> encoder.conv_norm_out + - encoder.quant_conv -> quant_conv (top-level) + - decoder.up.X -> decoder.up_blocks.(num_blocks-1-X) (reversed order!) + - decoder.post_quant_conv -> post_quant_conv (top-level) + - *.nin_shortcut -> *.conv_shortcut + """ + converted: dict = {} + num_up_blocks = 4 # Standard VAE has 4 up blocks + + for old_key, tensor in sd.items(): + new_key = old_key + + # Encoder down blocks: encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y + match = re.match(r"encoder\.down\.(\d+)\.block\.(\d+)\.(.*)", old_key) + if match: + block_idx, resnet_idx, rest = match.groups() + rest = rest.replace("nin_shortcut", "conv_shortcut") + new_key = f"encoder.down_blocks.{block_idx}.resnets.{resnet_idx}.{rest}" + converted[new_key] = tensor + continue + + # Encoder downsamplers: encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv + match = re.match(r"encoder\.down\.(\d+)\.downsample\.conv\.(.*)", old_key) + if match: + block_idx, rest = match.groups() + new_key = f"encoder.down_blocks.{block_idx}.downsamplers.0.conv.{rest}" + converted[new_key] = tensor + continue + + # Encoder mid block resnets: encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1 + match = re.match(r"encoder\.mid\.block_(\d+)\.(.*)", old_key) + if match: + block_num, rest = match.groups() + resnet_idx = int(block_num) - 1 # block_1 -> resnets.0, block_2 -> resnets.1 + new_key = f"encoder.mid_block.resnets.{resnet_idx}.{rest}" + converted[new_key] = tensor + continue + + # Encoder mid block attention: encoder.mid.attn_1.* -> encoder.mid_block.attentions.0.* + match = re.match(r"encoder\.mid\.attn_1\.(.*)", old_key) + if match: + new_key, tensor = _convert_flux2_vae_mid_attention_key(match.group(1), tensor, "encoder") + converted[new_key] = tensor + continue + + # Encoder norm_out -> conv_norm_out + if old_key.startswith("encoder.norm_out."): + new_key = old_key.replace("encoder.norm_out.", "encoder.conv_norm_out.") + converted[new_key] = tensor + continue + + # Encoder quant_conv -> quant_conv (move to top level) + if old_key.startswith("encoder.quant_conv."): + new_key = old_key.replace("encoder.quant_conv.", "quant_conv.") + converted[new_key] = tensor + continue + + # Decoder up blocks (reversed order!): decoder.up.X -> decoder.up_blocks.(num_blocks-1-X) + match = re.match(r"decoder\.up\.(\d+)\.block\.(\d+)\.(.*)", old_key) + if match: + block_idx, resnet_idx, rest = match.groups() + new_block_idx = num_up_blocks - 1 - int(block_idx) + rest = rest.replace("nin_shortcut", "conv_shortcut") + new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.{rest}" + converted[new_key] = tensor + continue + + # Decoder upsamplers (reversed order!) + match = re.match(r"decoder\.up\.(\d+)\.upsample\.conv\.(.*)", old_key) + if match: + block_idx, rest = match.groups() + new_block_idx = num_up_blocks - 1 - int(block_idx) + new_key = f"decoder.up_blocks.{new_block_idx}.upsamplers.0.conv.{rest}" + converted[new_key] = tensor + continue + + # Decoder mid block resnets: decoder.mid.block_1/2 -> decoder.mid_block.resnets.0/1 + match = re.match(r"decoder\.mid\.block_(\d+)\.(.*)", old_key) + if match: + block_num, rest = match.groups() + resnet_idx = int(block_num) - 1 + new_key = f"decoder.mid_block.resnets.{resnet_idx}.{rest}" + converted[new_key] = tensor + continue + + # Decoder mid block attention: decoder.mid.attn_1.* -> decoder.mid_block.attentions.0.* + match = re.match(r"decoder\.mid\.attn_1\.(.*)", old_key) + if match: + new_key, tensor = _convert_flux2_vae_mid_attention_key(match.group(1), tensor, "decoder") + converted[new_key] = tensor + continue + + # Decoder norm_out -> conv_norm_out + if old_key.startswith("decoder.norm_out."): + new_key = old_key.replace("decoder.norm_out.", "decoder.conv_norm_out.") + converted[new_key] = tensor + continue + + # Decoder post_quant_conv -> post_quant_conv (move to top level) + if old_key.startswith("decoder.post_quant_conv."): + new_key = old_key.replace("decoder.post_quant_conv.", "post_quant_conv.") + converted[new_key] = tensor + continue + + # Keep other keys as-is (like encoder.conv_in, decoder.conv_in, decoder.conv_out, bn.*) + converted[new_key] = tensor + + return converted diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py index 11cd497720b..90919794999 100644 --- a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py @@ -87,6 +87,43 @@ def _dequantize_comfyui_fp8(sd: dict, compute_dtype: torch.dtype) -> int: return count +def _remap_qwen_vl_checkpoint_keys(sd: dict) -> dict: + """Remap legacy ComfyUI Qwen2.5-VL single-file keys to the transformers layout. + + ComfyUI single-file checkpoints use the legacy Qwen2.5-VL key layout + (`visual.X`, `model.X`); transformers ≥4.50 expects `model.visual.X` and + `model.language_model.X`. This applies the same conversion mapping that + `Qwen2_5_VLForConditionalGeneration.from_pretrained` would, since + `load_state_dict` does not. + + transformers ≤4.x exposed this as `_checkpoint_conversion_mapping`, but 5.x + dropped it (returns `{}`), so we fall back to the legacy mapping ourselves. The + negative lookahead keeps already-converted keys untouched, so the remap is safe + (and idempotent) for both legacy and new-layout single-file checkpoints. + """ + import re + + from transformers import Qwen2_5_VLForConditionalGeneration + + key_mapping = Qwen2_5_VLForConditionalGeneration._checkpoint_conversion_mapping or { + r"^visual": "model.visual", + r"^model(?!\.(language_model|visual))": "model.language_model", + } + if not key_mapping: + return sd + + remapped_sd: dict = {} + for old_key, tensor in sd.items(): + new_key = old_key + if isinstance(old_key, str): + for pattern, replacement in key_mapping.items(): + new_key, n_replace = re.subn(pattern, replacement, new_key) + if n_replace > 0: + break + remapped_sd[new_key] = tensor + return remapped_sd + + def _strip_quantization_metadata(sd: dict) -> None: """Strip ComfyUI fp8 quantization metadata keys in-place.""" keys_to_drop = [ @@ -415,8 +452,6 @@ def _load_tokenizer_with_offline_fallback(self) -> AnyModel: ) from e def _load_text_encoder_from_singlefile(self, config: QwenVLEncoder_Checkpoint_Config) -> AnyModel: - import re - from safetensors.torch import load_file from transformers import AutoConfig, Qwen2_5_VLForConditionalGeneration @@ -440,29 +475,9 @@ def _load_text_encoder_from_singlefile(self, config: QwenVLEncoder_Checkpoint_Co _strip_quantization_metadata(sd) # ComfyUI single-file checkpoints use the legacy Qwen2.5-VL key layout - # (`visual.X`, `model.X`); transformers ≥4.50 expects `model.visual.X` and - # `model.language_model.X`. Apply the same conversion mapping that - # `Qwen2_5_VLForConditionalGeneration.from_pretrained` would, since - # `load_state_dict` does not. - # - # transformers ≤4.x exposed this as `_checkpoint_conversion_mapping`, but 5.x - # dropped it (returns `{}`), so we fall back to the legacy mapping ourselves. - # The negative lookahead keeps already-converted checkpoints untouched, so the - # remap is safe to apply to both legacy and new-layout single-file checkpoints. - key_mapping = Qwen2_5_VLForConditionalGeneration._checkpoint_conversion_mapping or { - r"^visual": "model.visual", - r"^model(?!\.(language_model|visual))": "model.language_model", - } - if key_mapping: - remapped_sd: dict[str, torch.Tensor] = {} - for old_key, tensor in sd.items(): - new_key = old_key - for pattern, replacement in key_mapping.items(): - new_key, n_replace = re.subn(pattern, replacement, new_key) - if n_replace > 0: - break - remapped_sd[new_key] = tensor - sd = remapped_sd + # (`visual.X`, `model.X`); remap to the `model.visual.X` / `model.language_model.X` + # layout transformers expects. See `_remap_qwen_vl_checkpoint_keys` for details. + sd = _remap_qwen_vl_checkpoint_keys(sd) # Cast to compute dtype (skip integer/index tensors) for k in list(sd.keys()): diff --git a/tests/backend/model_manager/load/state_dicts/anima_comfyui_keys.py b/tests/backend/model_manager/load/state_dicts/anima_comfyui_keys.py new file mode 100644 index 00000000000..addd8ca6fc2 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/anima_comfyui_keys.py @@ -0,0 +1,155 @@ +"""Representative key layout of an official Anima transformer single-file checkpoint. + +Captured from `anima-base-v1.0.safetensors`. The full checkpoint has 685 tensors under the +`net.` prefix; this fixture keeps `net.blocks.0.*` plus all non-block `net.*` keys. Used to +exercise `_strip_anima_bundle_prefix` (the `net.` -> unprefixed strip). The ComfyUI-bundled +`model.diffusion_model.*` variant is covered by a small synthetic dict in the test. +""" + +state_dict_keys: dict[str, list[int]] = { + "net.blocks.0.adaln_modulation_cross_attn.1.weight": [256, 2048], + "net.blocks.0.adaln_modulation_cross_attn.2.weight": [6144, 256], + "net.blocks.0.adaln_modulation_mlp.1.weight": [256, 2048], + "net.blocks.0.adaln_modulation_mlp.2.weight": [6144, 256], + "net.blocks.0.adaln_modulation_self_attn.1.weight": [256, 2048], + "net.blocks.0.adaln_modulation_self_attn.2.weight": [6144, 256], + "net.blocks.0.cross_attn.k_norm.weight": [128], + "net.blocks.0.cross_attn.k_proj.weight": [2048, 1024], + "net.blocks.0.cross_attn.output_proj.weight": [2048, 2048], + "net.blocks.0.cross_attn.q_norm.weight": [128], + "net.blocks.0.cross_attn.q_proj.weight": [2048, 2048], + "net.blocks.0.cross_attn.v_proj.weight": [2048, 1024], + "net.blocks.0.mlp.layer1.weight": [8192, 2048], + "net.blocks.0.mlp.layer2.weight": [2048, 8192], + "net.blocks.0.self_attn.k_norm.weight": [128], + "net.blocks.0.self_attn.k_proj.weight": [2048, 2048], + "net.blocks.0.self_attn.output_proj.weight": [2048, 2048], + "net.blocks.0.self_attn.q_norm.weight": [128], + "net.blocks.0.self_attn.q_proj.weight": [2048, 2048], + "net.blocks.0.self_attn.v_proj.weight": [2048, 2048], + "net.final_layer.adaln_modulation.1.weight": [256, 2048], + "net.final_layer.adaln_modulation.2.weight": [4096, 256], + "net.final_layer.linear.weight": [64, 2048], + "net.llm_adapter.blocks.0.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.0.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.0.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.mlp.0.bias": [4096], + "net.llm_adapter.blocks.0.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.0.mlp.2.bias": [1024], + "net.llm_adapter.blocks.0.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.0.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.0.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.0.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.0.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.0.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.0.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.0.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.1.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.1.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.mlp.0.bias": [4096], + "net.llm_adapter.blocks.1.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.1.mlp.2.bias": [1024], + "net.llm_adapter.blocks.1.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.1.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.1.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.1.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.1.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.1.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.1.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.1.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.2.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.2.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.mlp.0.bias": [4096], + "net.llm_adapter.blocks.2.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.2.mlp.2.bias": [1024], + "net.llm_adapter.blocks.2.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.2.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.2.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.2.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.2.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.2.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.2.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.2.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.3.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.3.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.mlp.0.bias": [4096], + "net.llm_adapter.blocks.3.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.3.mlp.2.bias": [1024], + "net.llm_adapter.blocks.3.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.3.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.3.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.3.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.3.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.3.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.3.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.3.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.4.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.4.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.mlp.0.bias": [4096], + "net.llm_adapter.blocks.4.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.4.mlp.2.bias": [1024], + "net.llm_adapter.blocks.4.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.4.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.4.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.4.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.4.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.4.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.4.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.4.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.cross_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.5.cross_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.cross_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.cross_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.5.cross_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.cross_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.mlp.0.bias": [4096], + "net.llm_adapter.blocks.5.mlp.0.weight": [4096, 1024], + "net.llm_adapter.blocks.5.mlp.2.bias": [1024], + "net.llm_adapter.blocks.5.mlp.2.weight": [1024, 4096], + "net.llm_adapter.blocks.5.norm_cross_attn.weight": [1024], + "net.llm_adapter.blocks.5.norm_mlp.weight": [1024], + "net.llm_adapter.blocks.5.norm_self_attn.weight": [1024], + "net.llm_adapter.blocks.5.self_attn.k_norm.weight": [64], + "net.llm_adapter.blocks.5.self_attn.k_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.self_attn.o_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.self_attn.q_norm.weight": [64], + "net.llm_adapter.blocks.5.self_attn.q_proj.weight": [1024, 1024], + "net.llm_adapter.blocks.5.self_attn.v_proj.weight": [1024, 1024], + "net.llm_adapter.embed.weight": [32128, 1024], + "net.llm_adapter.norm.weight": [1024], + "net.llm_adapter.out_proj.bias": [1024], + "net.llm_adapter.out_proj.weight": [1024, 1024], + "net.t_embedder.1.linear_1.weight": [2048, 2048], + "net.t_embedder.1.linear_2.weight": [6144, 2048], + "net.t_embedding_norm.weight": [2048], + "net.x_embedder.proj.1.weight": [2048, 68], +} diff --git a/tests/backend/model_manager/load/state_dicts/flux2_transformer_bfl_keys.py b/tests/backend/model_manager/load/state_dicts/flux2_transformer_bfl_keys.py new file mode 100644 index 00000000000..627ec9868b6 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/flux2_transformer_bfl_keys.py @@ -0,0 +1,41 @@ +"""Representative BFL-format key layout of a FLUX.2 transformer single-file checkpoint. + +Captured from `flux-2-klein-9b-kv.safetensors` (FLUX.2 Klein 9B, bf16). The full +checkpoint has 201 tensors (8 double blocks, 24 single blocks); this +fixture keeps every top-level key plus block 0 of the double/single stacks, which is enough +to exercise `convert_flux2_bfl_to_diffusers` (fused-QKV split, block renames, adaLN +scale/shift swap) and validate against a single-layer `Flux2Transformer2DModel`. + +BFL layout uses `double_blocks.*`, `single_blocks.*`, `img_in`, `txt_in`, `time_in`, +`*_modulation.lin`, `final_layer.*`; diffusers expects `transformer_blocks.*`, +`single_transformer_blocks.*`, `x_embedder`, `context_embedder`, `time_guidance_embed.*`, +`proj_out`, `norm_out`. +""" + +state_dict_keys: dict[str, list[int]] = { + "double_blocks.0.img_attn.norm.key_norm.scale": [128], + "double_blocks.0.img_attn.norm.query_norm.scale": [128], + "double_blocks.0.img_attn.proj.weight": [4096, 4096], + "double_blocks.0.img_attn.qkv.weight": [12288, 4096], + "double_blocks.0.img_mlp.0.weight": [24576, 4096], + "double_blocks.0.img_mlp.2.weight": [4096, 12288], + "double_blocks.0.txt_attn.norm.key_norm.scale": [128], + "double_blocks.0.txt_attn.norm.query_norm.scale": [128], + "double_blocks.0.txt_attn.proj.weight": [4096, 4096], + "double_blocks.0.txt_attn.qkv.weight": [12288, 4096], + "double_blocks.0.txt_mlp.0.weight": [24576, 4096], + "double_blocks.0.txt_mlp.2.weight": [4096, 12288], + "double_stream_modulation_img.lin.weight": [24576, 4096], + "double_stream_modulation_txt.lin.weight": [24576, 4096], + "final_layer.adaLN_modulation.1.weight": [8192, 4096], + "final_layer.linear.weight": [128, 4096], + "img_in.weight": [4096, 128], + "single_blocks.0.linear1.weight": [36864, 4096], + "single_blocks.0.linear2.weight": [4096, 16384], + "single_blocks.0.norm.key_norm.scale": [128], + "single_blocks.0.norm.query_norm.scale": [128], + "single_stream_modulation.lin.weight": [12288, 4096], + "time_in.in_layer.weight": [4096, 256], + "time_in.out_layer.weight": [4096, 4096], + "txt_in.weight": [4096, 12288], +} diff --git a/tests/backend/model_manager/load/state_dicts/flux2_vae_bfl_keys.py b/tests/backend/model_manager/load/state_dicts/flux2_vae_bfl_keys.py new file mode 100644 index 00000000000..8d1e7e45db6 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/flux2_vae_bfl_keys.py @@ -0,0 +1,266 @@ +"""BFL-format key layout of a FLUX.2 VAE single-file checkpoint (full, 251 keys). + +Captured from `flux2-vae.safetensors` (standard FLUX.2 VAE, block_out_channels=(128,256,512,512)). +The full key set is kept (the VAE is small), so `convert_flux2_vae_bfl_to_diffusers` can be +validated for *complete* coverage against `AutoencoderKLFlux2` — every converted key must be a +real parameter and every parameter must be covered. + +BFL layout uses `encoder.down.*`, `decoder.up.*` (reversed order!), `{enc,dec}.mid.block_N`, +`{enc,dec}.mid.attn_1.*`, `norm_out`, `encoder.quant_conv`, `decoder.post_quant_conv`; +diffusers expects `down_blocks`/`up_blocks`/`mid_block.resnets`/`mid_block.attentions`/ +`conv_norm_out` and top-level `quant_conv`/`post_quant_conv`. +""" + +state_dict_keys: dict[str, list[int]] = { + "bn.num_batches_tracked": [], + "bn.running_mean": [128], + "bn.running_var": [128], + "decoder.conv_in.bias": [512], + "decoder.conv_in.weight": [512, 32, 3, 3], + "decoder.conv_norm_out.bias": [128], + "decoder.conv_norm_out.weight": [128], + "decoder.conv_out.bias": [3], + "decoder.conv_out.weight": [3, 128, 3, 3], + "decoder.mid_block.attentions.0.group_norm.bias": [512], + "decoder.mid_block.attentions.0.group_norm.weight": [512], + "decoder.mid_block.attentions.0.to_k.bias": [512], + "decoder.mid_block.attentions.0.to_k.weight": [512, 512], + "decoder.mid_block.attentions.0.to_out.0.bias": [512], + "decoder.mid_block.attentions.0.to_out.0.weight": [512, 512], + "decoder.mid_block.attentions.0.to_q.bias": [512], + "decoder.mid_block.attentions.0.to_q.weight": [512, 512], + "decoder.mid_block.attentions.0.to_v.bias": [512], + "decoder.mid_block.attentions.0.to_v.weight": [512, 512], + "decoder.mid_block.resnets.0.conv1.bias": [512], + "decoder.mid_block.resnets.0.conv1.weight": [512, 512, 3, 3], + "decoder.mid_block.resnets.0.conv2.bias": [512], + "decoder.mid_block.resnets.0.conv2.weight": [512, 512, 3, 3], + "decoder.mid_block.resnets.0.norm1.bias": [512], + "decoder.mid_block.resnets.0.norm1.weight": [512], + "decoder.mid_block.resnets.0.norm2.bias": [512], + "decoder.mid_block.resnets.0.norm2.weight": [512], + "decoder.mid_block.resnets.1.conv1.bias": [512], + "decoder.mid_block.resnets.1.conv1.weight": [512, 512, 3, 3], + "decoder.mid_block.resnets.1.conv2.bias": [512], + "decoder.mid_block.resnets.1.conv2.weight": [512, 512, 3, 3], + "decoder.mid_block.resnets.1.norm1.bias": [512], + "decoder.mid_block.resnets.1.norm1.weight": [512], + "decoder.mid_block.resnets.1.norm2.bias": [512], + "decoder.mid_block.resnets.1.norm2.weight": [512], + "decoder.up_blocks.0.resnets.0.conv1.bias": [512], + "decoder.up_blocks.0.resnets.0.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.0.conv2.bias": [512], + "decoder.up_blocks.0.resnets.0.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.0.norm1.bias": [512], + "decoder.up_blocks.0.resnets.0.norm1.weight": [512], + "decoder.up_blocks.0.resnets.0.norm2.bias": [512], + "decoder.up_blocks.0.resnets.0.norm2.weight": [512], + "decoder.up_blocks.0.resnets.1.conv1.bias": [512], + "decoder.up_blocks.0.resnets.1.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.1.conv2.bias": [512], + "decoder.up_blocks.0.resnets.1.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.1.norm1.bias": [512], + "decoder.up_blocks.0.resnets.1.norm1.weight": [512], + "decoder.up_blocks.0.resnets.1.norm2.bias": [512], + "decoder.up_blocks.0.resnets.1.norm2.weight": [512], + "decoder.up_blocks.0.resnets.2.conv1.bias": [512], + "decoder.up_blocks.0.resnets.2.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.2.conv2.bias": [512], + "decoder.up_blocks.0.resnets.2.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.0.resnets.2.norm1.bias": [512], + "decoder.up_blocks.0.resnets.2.norm1.weight": [512], + "decoder.up_blocks.0.resnets.2.norm2.bias": [512], + "decoder.up_blocks.0.resnets.2.norm2.weight": [512], + "decoder.up_blocks.0.upsamplers.0.conv.bias": [512], + "decoder.up_blocks.0.upsamplers.0.conv.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.0.conv1.bias": [512], + "decoder.up_blocks.1.resnets.0.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.0.conv2.bias": [512], + "decoder.up_blocks.1.resnets.0.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.0.norm1.bias": [512], + "decoder.up_blocks.1.resnets.0.norm1.weight": [512], + "decoder.up_blocks.1.resnets.0.norm2.bias": [512], + "decoder.up_blocks.1.resnets.0.norm2.weight": [512], + "decoder.up_blocks.1.resnets.1.conv1.bias": [512], + "decoder.up_blocks.1.resnets.1.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.1.conv2.bias": [512], + "decoder.up_blocks.1.resnets.1.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.1.norm1.bias": [512], + "decoder.up_blocks.1.resnets.1.norm1.weight": [512], + "decoder.up_blocks.1.resnets.1.norm2.bias": [512], + "decoder.up_blocks.1.resnets.1.norm2.weight": [512], + "decoder.up_blocks.1.resnets.2.conv1.bias": [512], + "decoder.up_blocks.1.resnets.2.conv1.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.2.conv2.bias": [512], + "decoder.up_blocks.1.resnets.2.conv2.weight": [512, 512, 3, 3], + "decoder.up_blocks.1.resnets.2.norm1.bias": [512], + "decoder.up_blocks.1.resnets.2.norm1.weight": [512], + "decoder.up_blocks.1.resnets.2.norm2.bias": [512], + "decoder.up_blocks.1.resnets.2.norm2.weight": [512], + "decoder.up_blocks.1.upsamplers.0.conv.bias": [512], + "decoder.up_blocks.1.upsamplers.0.conv.weight": [512, 512, 3, 3], + "decoder.up_blocks.2.resnets.0.conv1.bias": [256], + "decoder.up_blocks.2.resnets.0.conv1.weight": [256, 512, 3, 3], + "decoder.up_blocks.2.resnets.0.conv2.bias": [256], + "decoder.up_blocks.2.resnets.0.conv2.weight": [256, 256, 3, 3], + "decoder.up_blocks.2.resnets.0.conv_shortcut.bias": [256], + "decoder.up_blocks.2.resnets.0.conv_shortcut.weight": [256, 512, 1, 1], + "decoder.up_blocks.2.resnets.0.norm1.bias": [512], + "decoder.up_blocks.2.resnets.0.norm1.weight": [512], + "decoder.up_blocks.2.resnets.0.norm2.bias": [256], + "decoder.up_blocks.2.resnets.0.norm2.weight": [256], + "decoder.up_blocks.2.resnets.1.conv1.bias": [256], + "decoder.up_blocks.2.resnets.1.conv1.weight": [256, 256, 3, 3], + "decoder.up_blocks.2.resnets.1.conv2.bias": [256], + "decoder.up_blocks.2.resnets.1.conv2.weight": [256, 256, 3, 3], + "decoder.up_blocks.2.resnets.1.norm1.bias": [256], + "decoder.up_blocks.2.resnets.1.norm1.weight": [256], + "decoder.up_blocks.2.resnets.1.norm2.bias": [256], + "decoder.up_blocks.2.resnets.1.norm2.weight": [256], + "decoder.up_blocks.2.resnets.2.conv1.bias": [256], + "decoder.up_blocks.2.resnets.2.conv1.weight": [256, 256, 3, 3], + "decoder.up_blocks.2.resnets.2.conv2.bias": [256], + "decoder.up_blocks.2.resnets.2.conv2.weight": [256, 256, 3, 3], + "decoder.up_blocks.2.resnets.2.norm1.bias": [256], + "decoder.up_blocks.2.resnets.2.norm1.weight": [256], + "decoder.up_blocks.2.resnets.2.norm2.bias": [256], + "decoder.up_blocks.2.resnets.2.norm2.weight": [256], + "decoder.up_blocks.2.upsamplers.0.conv.bias": [256], + "decoder.up_blocks.2.upsamplers.0.conv.weight": [256, 256, 3, 3], + "decoder.up_blocks.3.resnets.0.conv1.bias": [128], + "decoder.up_blocks.3.resnets.0.conv1.weight": [128, 256, 3, 3], + "decoder.up_blocks.3.resnets.0.conv2.bias": [128], + "decoder.up_blocks.3.resnets.0.conv2.weight": [128, 128, 3, 3], + "decoder.up_blocks.3.resnets.0.conv_shortcut.bias": [128], + "decoder.up_blocks.3.resnets.0.conv_shortcut.weight": [128, 256, 1, 1], + "decoder.up_blocks.3.resnets.0.norm1.bias": [256], + "decoder.up_blocks.3.resnets.0.norm1.weight": [256], + "decoder.up_blocks.3.resnets.0.norm2.bias": [128], + "decoder.up_blocks.3.resnets.0.norm2.weight": [128], + "decoder.up_blocks.3.resnets.1.conv1.bias": [128], + "decoder.up_blocks.3.resnets.1.conv1.weight": [128, 128, 3, 3], + "decoder.up_blocks.3.resnets.1.conv2.bias": [128], + "decoder.up_blocks.3.resnets.1.conv2.weight": [128, 128, 3, 3], + "decoder.up_blocks.3.resnets.1.norm1.bias": [128], + "decoder.up_blocks.3.resnets.1.norm1.weight": [128], + "decoder.up_blocks.3.resnets.1.norm2.bias": [128], + "decoder.up_blocks.3.resnets.1.norm2.weight": [128], + "decoder.up_blocks.3.resnets.2.conv1.bias": [128], + "decoder.up_blocks.3.resnets.2.conv1.weight": [128, 128, 3, 3], + "decoder.up_blocks.3.resnets.2.conv2.bias": [128], + "decoder.up_blocks.3.resnets.2.conv2.weight": [128, 128, 3, 3], + "decoder.up_blocks.3.resnets.2.norm1.bias": [128], + "decoder.up_blocks.3.resnets.2.norm1.weight": [128], + "decoder.up_blocks.3.resnets.2.norm2.bias": [128], + "decoder.up_blocks.3.resnets.2.norm2.weight": [128], + "encoder.conv_in.bias": [128], + "encoder.conv_in.weight": [128, 3, 3, 3], + "encoder.conv_norm_out.bias": [512], + "encoder.conv_norm_out.weight": [512], + "encoder.conv_out.bias": [64], + "encoder.conv_out.weight": [64, 512, 3, 3], + "encoder.down_blocks.0.downsamplers.0.conv.bias": [128], + "encoder.down_blocks.0.downsamplers.0.conv.weight": [128, 128, 3, 3], + "encoder.down_blocks.0.resnets.0.conv1.bias": [128], + "encoder.down_blocks.0.resnets.0.conv1.weight": [128, 128, 3, 3], + "encoder.down_blocks.0.resnets.0.conv2.bias": [128], + "encoder.down_blocks.0.resnets.0.conv2.weight": [128, 128, 3, 3], + "encoder.down_blocks.0.resnets.0.norm1.bias": [128], + "encoder.down_blocks.0.resnets.0.norm1.weight": [128], + "encoder.down_blocks.0.resnets.0.norm2.bias": [128], + "encoder.down_blocks.0.resnets.0.norm2.weight": [128], + "encoder.down_blocks.0.resnets.1.conv1.bias": [128], + "encoder.down_blocks.0.resnets.1.conv1.weight": [128, 128, 3, 3], + "encoder.down_blocks.0.resnets.1.conv2.bias": [128], + "encoder.down_blocks.0.resnets.1.conv2.weight": [128, 128, 3, 3], + "encoder.down_blocks.0.resnets.1.norm1.bias": [128], + "encoder.down_blocks.0.resnets.1.norm1.weight": [128], + "encoder.down_blocks.0.resnets.1.norm2.bias": [128], + "encoder.down_blocks.0.resnets.1.norm2.weight": [128], + "encoder.down_blocks.1.downsamplers.0.conv.bias": [256], + "encoder.down_blocks.1.downsamplers.0.conv.weight": [256, 256, 3, 3], + "encoder.down_blocks.1.resnets.0.conv1.bias": [256], + "encoder.down_blocks.1.resnets.0.conv1.weight": [256, 128, 3, 3], + "encoder.down_blocks.1.resnets.0.conv2.bias": [256], + "encoder.down_blocks.1.resnets.0.conv2.weight": [256, 256, 3, 3], + "encoder.down_blocks.1.resnets.0.conv_shortcut.bias": [256], + "encoder.down_blocks.1.resnets.0.conv_shortcut.weight": [256, 128, 1, 1], + "encoder.down_blocks.1.resnets.0.norm1.bias": [128], + "encoder.down_blocks.1.resnets.0.norm1.weight": [128], + "encoder.down_blocks.1.resnets.0.norm2.bias": [256], + "encoder.down_blocks.1.resnets.0.norm2.weight": [256], + "encoder.down_blocks.1.resnets.1.conv1.bias": [256], + "encoder.down_blocks.1.resnets.1.conv1.weight": [256, 256, 3, 3], + "encoder.down_blocks.1.resnets.1.conv2.bias": [256], + "encoder.down_blocks.1.resnets.1.conv2.weight": [256, 256, 3, 3], + "encoder.down_blocks.1.resnets.1.norm1.bias": [256], + "encoder.down_blocks.1.resnets.1.norm1.weight": [256], + "encoder.down_blocks.1.resnets.1.norm2.bias": [256], + "encoder.down_blocks.1.resnets.1.norm2.weight": [256], + "encoder.down_blocks.2.downsamplers.0.conv.bias": [512], + "encoder.down_blocks.2.downsamplers.0.conv.weight": [512, 512, 3, 3], + "encoder.down_blocks.2.resnets.0.conv1.bias": [512], + "encoder.down_blocks.2.resnets.0.conv1.weight": [512, 256, 3, 3], + "encoder.down_blocks.2.resnets.0.conv2.bias": [512], + "encoder.down_blocks.2.resnets.0.conv2.weight": [512, 512, 3, 3], + "encoder.down_blocks.2.resnets.0.conv_shortcut.bias": [512], + "encoder.down_blocks.2.resnets.0.conv_shortcut.weight": [512, 256, 1, 1], + "encoder.down_blocks.2.resnets.0.norm1.bias": [256], + "encoder.down_blocks.2.resnets.0.norm1.weight": [256], + "encoder.down_blocks.2.resnets.0.norm2.bias": [512], + "encoder.down_blocks.2.resnets.0.norm2.weight": [512], + "encoder.down_blocks.2.resnets.1.conv1.bias": [512], + "encoder.down_blocks.2.resnets.1.conv1.weight": [512, 512, 3, 3], + "encoder.down_blocks.2.resnets.1.conv2.bias": [512], + "encoder.down_blocks.2.resnets.1.conv2.weight": [512, 512, 3, 3], + "encoder.down_blocks.2.resnets.1.norm1.bias": [512], + "encoder.down_blocks.2.resnets.1.norm1.weight": [512], + "encoder.down_blocks.2.resnets.1.norm2.bias": [512], + "encoder.down_blocks.2.resnets.1.norm2.weight": [512], + "encoder.down_blocks.3.resnets.0.conv1.bias": [512], + "encoder.down_blocks.3.resnets.0.conv1.weight": [512, 512, 3, 3], + "encoder.down_blocks.3.resnets.0.conv2.bias": [512], + "encoder.down_blocks.3.resnets.0.conv2.weight": [512, 512, 3, 3], + "encoder.down_blocks.3.resnets.0.norm1.bias": [512], + "encoder.down_blocks.3.resnets.0.norm1.weight": [512], + "encoder.down_blocks.3.resnets.0.norm2.bias": [512], + "encoder.down_blocks.3.resnets.0.norm2.weight": [512], + "encoder.down_blocks.3.resnets.1.conv1.bias": [512], + "encoder.down_blocks.3.resnets.1.conv1.weight": [512, 512, 3, 3], + "encoder.down_blocks.3.resnets.1.conv2.bias": [512], + "encoder.down_blocks.3.resnets.1.conv2.weight": [512, 512, 3, 3], + "encoder.down_blocks.3.resnets.1.norm1.bias": [512], + "encoder.down_blocks.3.resnets.1.norm1.weight": [512], + "encoder.down_blocks.3.resnets.1.norm2.bias": [512], + "encoder.down_blocks.3.resnets.1.norm2.weight": [512], + "encoder.mid_block.attentions.0.group_norm.bias": [512], + "encoder.mid_block.attentions.0.group_norm.weight": [512], + "encoder.mid_block.attentions.0.to_k.bias": [512], + "encoder.mid_block.attentions.0.to_k.weight": [512, 512], + "encoder.mid_block.attentions.0.to_out.0.bias": [512], + "encoder.mid_block.attentions.0.to_out.0.weight": [512, 512], + "encoder.mid_block.attentions.0.to_q.bias": [512], + "encoder.mid_block.attentions.0.to_q.weight": [512, 512], + "encoder.mid_block.attentions.0.to_v.bias": [512], + "encoder.mid_block.attentions.0.to_v.weight": [512, 512], + "encoder.mid_block.resnets.0.conv1.bias": [512], + "encoder.mid_block.resnets.0.conv1.weight": [512, 512, 3, 3], + "encoder.mid_block.resnets.0.conv2.bias": [512], + "encoder.mid_block.resnets.0.conv2.weight": [512, 512, 3, 3], + "encoder.mid_block.resnets.0.norm1.bias": [512], + "encoder.mid_block.resnets.0.norm1.weight": [512], + "encoder.mid_block.resnets.0.norm2.bias": [512], + "encoder.mid_block.resnets.0.norm2.weight": [512], + "encoder.mid_block.resnets.1.conv1.bias": [512], + "encoder.mid_block.resnets.1.conv1.weight": [512, 512, 3, 3], + "encoder.mid_block.resnets.1.conv2.bias": [512], + "encoder.mid_block.resnets.1.conv2.weight": [512, 512, 3, 3], + "encoder.mid_block.resnets.1.norm1.bias": [512], + "encoder.mid_block.resnets.1.norm1.weight": [512], + "encoder.mid_block.resnets.1.norm2.bias": [512], + "encoder.mid_block.resnets.1.norm2.weight": [512], + "post_quant_conv.bias": [32], + "post_quant_conv.weight": [32, 32, 1, 1], + "quant_conv.bias": [64], + "quant_conv.weight": [64, 64, 1, 1], +} diff --git a/tests/backend/model_manager/load/state_dicts/qwen_vl_encoder_comfyui_keys.py b/tests/backend/model_manager/load/state_dicts/qwen_vl_encoder_comfyui_keys.py new file mode 100644 index 00000000000..d139ef56360 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/qwen_vl_encoder_comfyui_keys.py @@ -0,0 +1,77 @@ +"""Representative key layout of a ComfyUI single-file Qwen2.5-VL encoder checkpoint. + +Captured from `qwen_2.5_vl_7b_fp8_scaled.safetensors` (Qwen2.5-VL-7B, ComfyUI fp8_scaled). +The full checkpoint has 1446 tensors (32 visual blocks, 28 language layers); this +fixture keeps every top-level/structural key plus block 0 of each repeated stack, which is +enough to exercise the `visual.* -> model.visual.*` / `model.* -> model.language_model.*` +remap and the fp8 metadata stripping without shipping a ~1446-key dict. + +Legacy ComfyUI layout uses `visual.*`, `model.*`, `lm_head.*` (transformers >=4.50 expects +`model.visual.*` and `model.language_model.*`). `scale_weight` / `scale_input` / `scaled_fp8` +are ComfyUI fp8 quantization metadata. +""" + +state_dict_keys: dict[str, list[int]] = { + "visual.blocks.0.attn.proj.bias": [1280], + "visual.blocks.0.attn.proj.scale_input": [], + "visual.blocks.0.attn.proj.scale_weight": [], + "visual.blocks.0.attn.proj.weight": [1280, 1280], + "visual.blocks.0.attn.qkv.bias": [3840], + "visual.blocks.0.attn.qkv.scale_input": [], + "visual.blocks.0.attn.qkv.scale_weight": [], + "visual.blocks.0.attn.qkv.weight": [3840, 1280], + "visual.blocks.0.mlp.down_proj.bias": [1280], + "visual.blocks.0.mlp.down_proj.scale_input": [], + "visual.blocks.0.mlp.down_proj.scale_weight": [], + "visual.blocks.0.mlp.down_proj.weight": [1280, 3420], + "visual.blocks.0.mlp.gate_proj.bias": [3420], + "visual.blocks.0.mlp.gate_proj.scale_input": [], + "visual.blocks.0.mlp.gate_proj.scale_weight": [], + "visual.blocks.0.mlp.gate_proj.weight": [3420, 1280], + "visual.blocks.0.mlp.up_proj.bias": [3420], + "visual.blocks.0.mlp.up_proj.scale_input": [], + "visual.blocks.0.mlp.up_proj.scale_weight": [], + "visual.blocks.0.mlp.up_proj.weight": [3420, 1280], + "visual.blocks.0.norm1.weight": [1280], + "visual.blocks.0.norm2.weight": [1280], + "visual.merger.ln_q.weight": [1280], + "visual.merger.mlp.0.bias": [5120], + "visual.merger.mlp.0.scale_input": [], + "visual.merger.mlp.0.scale_weight": [], + "visual.merger.mlp.0.weight": [5120, 5120], + "visual.merger.mlp.2.bias": [3584], + "visual.merger.mlp.2.scale_input": [], + "visual.merger.mlp.2.scale_weight": [], + "visual.merger.mlp.2.weight": [3584, 5120], + "visual.patch_embed.proj.weight": [1280, 3, 2, 14, 14], + "model.embed_tokens.weight": [152064, 3584], + "model.layers.0.input_layernorm.weight": [3584], + "model.layers.0.mlp.down_proj.scale_input": [], + "model.layers.0.mlp.down_proj.scale_weight": [], + "model.layers.0.mlp.down_proj.weight": [3584, 18944], + "model.layers.0.mlp.gate_proj.scale_input": [], + "model.layers.0.mlp.gate_proj.scale_weight": [], + "model.layers.0.mlp.gate_proj.weight": [18944, 3584], + "model.layers.0.mlp.up_proj.scale_input": [], + "model.layers.0.mlp.up_proj.scale_weight": [], + "model.layers.0.mlp.up_proj.weight": [18944, 3584], + "model.layers.0.post_attention_layernorm.weight": [3584], + "model.layers.0.self_attn.k_proj.bias": [512], + "model.layers.0.self_attn.k_proj.scale_input": [], + "model.layers.0.self_attn.k_proj.scale_weight": [], + "model.layers.0.self_attn.k_proj.weight": [512, 3584], + "model.layers.0.self_attn.o_proj.scale_input": [], + "model.layers.0.self_attn.o_proj.scale_weight": [], + "model.layers.0.self_attn.o_proj.weight": [3584, 3584], + "model.layers.0.self_attn.q_proj.bias": [3584], + "model.layers.0.self_attn.q_proj.scale_input": [], + "model.layers.0.self_attn.q_proj.scale_weight": [], + "model.layers.0.self_attn.q_proj.weight": [3584, 3584], + "model.layers.0.self_attn.v_proj.bias": [512], + "model.layers.0.self_attn.v_proj.scale_input": [], + "model.layers.0.self_attn.v_proj.scale_weight": [], + "model.layers.0.self_attn.v_proj.weight": [512, 3584], + "model.norm.weight": [3584], + "lm_head.weight": [152064, 3584], + "scaled_fp8": [0], +} diff --git a/tests/backend/model_manager/load/state_dicts/utils.py b/tests/backend/model_manager/load/state_dicts/utils.py new file mode 100644 index 00000000000..f751490f196 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/utils.py @@ -0,0 +1,13 @@ +"""Shared helpers for model-loader state-dict fixtures. + +Mirrors `tests/backend/patches/lora_conversions/lora_state_dicts/utils.py`: a fixture module +exports `state_dict_keys: dict[str, list[int]]` (key name -> shape, captured from a real +checkpoint) and tests expand it to a mock state dict with `keys_to_mock_state_dict()`. +""" + +import torch + + +def keys_to_mock_state_dict(keys: dict[str, list[int]]) -> dict[str, torch.Tensor]: + """Build a state dict of empty tensors from a {key: shape} mapping.""" + return {k: torch.empty(shape) for k, shape in keys.items()} diff --git a/tests/backend/model_manager/load/state_dicts/z_image_transformer_comfyui_keys.py b/tests/backend/model_manager/load/state_dicts/z_image_transformer_comfyui_keys.py new file mode 100644 index 00000000000..56cce5381ed --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/z_image_transformer_comfyui_keys.py @@ -0,0 +1,69 @@ +"""Representative ComfyUI key layout of a Z-Image transformer single-file checkpoint. + +Captured from `zimageTurboBadmilk_v10.safetensors` (Z-Image Turbo) after stripping the +`model.diffusion_model.` prefix -- this is exactly the input `_convert_z_image_gguf_to_diffusers` +receives (the converter runs on both the checkpoint and GGUF paths after prefix stripping). +The full transformer has 453 keys (context_refiner / noise_refiner / layers stacks); this +fixture keeps block 0 of each stack plus all non-block keys, and adds a synthetic +`norm_final.weight` to exercise the skip branch. + +Legacy layout uses fused `*.attention.qkv.*`, `*.attention.out.*`, `*.attention.q_norm/k_norm`, +`x_embedder.*`, `final_layer.*`, `x_pad_token`, `cap_pad_token`; diffusers expects split +`to_q/to_k/to_v`, `to_out.0`, `norm_q/norm_k`, `all_x_embedder.2-1.*`, `all_final_layer.2-1.*`. +""" + +state_dict_keys: dict[str, list[int]] = { + "cap_embedder.0.weight": [2560], + "cap_embedder.1.bias": [3840], + "cap_embedder.1.weight": [3840, 2560], + "cap_pad_token": [1, 3840], + "context_refiner.0.attention.k_norm.weight": [128], + "context_refiner.0.attention.out.weight": [3840, 3840], + "context_refiner.0.attention.q_norm.weight": [128], + "context_refiner.0.attention.qkv.weight": [11520, 3840], + "context_refiner.0.attention_norm1.weight": [3840], + "context_refiner.0.attention_norm2.weight": [3840], + "context_refiner.0.feed_forward.w1.weight": [10240, 3840], + "context_refiner.0.feed_forward.w2.weight": [3840, 10240], + "context_refiner.0.feed_forward.w3.weight": [10240, 3840], + "context_refiner.0.ffn_norm1.weight": [3840], + "context_refiner.0.ffn_norm2.weight": [3840], + "final_layer.adaLN_modulation.1.bias": [3840], + "final_layer.adaLN_modulation.1.weight": [3840, 256], + "final_layer.linear.bias": [64], + "final_layer.linear.weight": [64, 3840], + "layers.0.adaLN_modulation.0.bias": [15360], + "layers.0.adaLN_modulation.0.weight": [15360, 256], + "layers.0.attention.k_norm.weight": [128], + "layers.0.attention.out.weight": [3840, 3840], + "layers.0.attention.q_norm.weight": [128], + "layers.0.attention.qkv.weight": [11520, 3840], + "layers.0.attention_norm1.weight": [3840], + "layers.0.attention_norm2.weight": [3840], + "layers.0.feed_forward.w1.weight": [10240, 3840], + "layers.0.feed_forward.w2.weight": [3840, 10240], + "layers.0.feed_forward.w3.weight": [10240, 3840], + "layers.0.ffn_norm1.weight": [3840], + "layers.0.ffn_norm2.weight": [3840], + "noise_refiner.0.adaLN_modulation.0.bias": [15360], + "noise_refiner.0.adaLN_modulation.0.weight": [15360, 256], + "noise_refiner.0.attention.k_norm.weight": [128], + "noise_refiner.0.attention.out.weight": [3840, 3840], + "noise_refiner.0.attention.q_norm.weight": [128], + "noise_refiner.0.attention.qkv.weight": [11520, 3840], + "noise_refiner.0.attention_norm1.weight": [3840], + "noise_refiner.0.attention_norm2.weight": [3840], + "noise_refiner.0.feed_forward.w1.weight": [10240, 3840], + "noise_refiner.0.feed_forward.w2.weight": [3840, 10240], + "noise_refiner.0.feed_forward.w3.weight": [10240, 3840], + "noise_refiner.0.ffn_norm1.weight": [3840], + "noise_refiner.0.ffn_norm2.weight": [3840], + "t_embedder.mlp.0.bias": [1024], + "t_embedder.mlp.0.weight": [1024, 256], + "t_embedder.mlp.2.bias": [256], + "t_embedder.mlp.2.weight": [256, 1024], + "x_embedder.bias": [3840], + "x_embedder.weight": [3840, 64], + "x_pad_token": [1, 3840], + "norm_final.weight": [2304], +} diff --git a/tests/backend/model_manager/load/test_anima_state_dict_utils.py b/tests/backend/model_manager/load/test_anima_state_dict_utils.py new file mode 100644 index 00000000000..769ea67df88 --- /dev/null +++ b/tests/backend/model_manager/load/test_anima_state_dict_utils.py @@ -0,0 +1,38 @@ +"""Unit tests for the Anima single-file prefix-stripping helper.""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.anima import _strip_anima_bundle_prefix +from tests.backend.model_manager.load.state_dicts.anima_comfyui_keys import state_dict_keys as anima_keys +from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict + + +class TestStripAnimaBundlePrefix: + def test_official_net_prefix_is_stripped(self): + sd = keys_to_mock_state_dict(anima_keys) + assert all(k.startswith("net.") for k in sd) + + out = _strip_anima_bundle_prefix(sd) + + assert len(out) == len(sd) + assert not any(k.startswith("net.") for k in out) + # Every key had exactly its `net.` prefix removed. + assert {"net." + k for k in out} == set(sd.keys()) + + def test_comfyui_bundle_keeps_only_transformer_keys(self): + # ComfyUI bundles the transformer under `model.diffusion_model.` alongside the VAE and + # text encoder, which must be dropped. + sd = { + "model.diffusion_model.blocks.0.attn.qkv.weight": torch.empty(1), + "model.diffusion_model.final_layer.weight": torch.empty(1), + "first_stage_model.encoder.conv_in.weight": torch.empty(1), + "cond_stage_model.transformer.embeddings.weight": torch.empty(1), + } + + out = _strip_anima_bundle_prefix(sd) + + assert set(out.keys()) == {"blocks.0.attn.qkv.weight", "final_layer.weight"} + + def test_no_known_prefix_is_a_noop(self): + sd = {"blocks.0.attn.qkv.weight": torch.empty(1)} + assert _strip_anima_bundle_prefix(sd) is sd diff --git a/tests/backend/model_manager/load/test_flux2_state_dict_utils.py b/tests/backend/model_manager/load/test_flux2_state_dict_utils.py new file mode 100644 index 00000000000..8425118af11 --- /dev/null +++ b/tests/backend/model_manager/load/test_flux2_state_dict_utils.py @@ -0,0 +1,101 @@ +"""Unit tests for the FLUX.2 BFL->diffusers state-dict converters. + +Fixtures are captured from real single-file checkpoints (see the fixture module docstrings). +The meta-device tests instantiate the actual diffusers architectures with `init_empty_weights` +(no real weights, no GPU) and assert that every converted key is a real parameter -- the same +kind of check that would have caught the Qwen VL remap regression. +""" + +import accelerate +import torch + +from invokeai.backend.model_manager.load.model_loaders.flux2_state_dict_utils import ( + _flux2_swap_scale_shift, + convert_flux2_bfl_to_diffusers, + convert_flux2_vae_bfl_to_diffusers, +) +from tests.backend.model_manager.load.state_dicts.flux2_transformer_bfl_keys import ( + state_dict_keys as flux2_transformer_keys, +) +from tests.backend.model_manager.load.state_dicts.flux2_vae_bfl_keys import ( + state_dict_keys as flux2_vae_keys, +) +from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict + + +class TestConvertFlux2Transformer: + def test_fused_qkv_is_split_and_blocks_renamed(self): + sd = keys_to_mock_state_dict(flux2_transformer_keys) + + converted = convert_flux2_bfl_to_diffusers(sd) + + # Fused img/txt QKV are split into separate projections. + assert "transformer_blocks.0.attn.to_q.weight" in converted + assert "transformer_blocks.0.attn.to_k.weight" in converted + assert "transformer_blocks.0.attn.to_v.weight" in converted + assert "transformer_blocks.0.attn.add_q_proj.weight" in converted + # No fused/BFL-named keys remain. + assert not any("img_attn.qkv" in k or "double_blocks." in k or "single_blocks." in k for k in converted) + # Top-level renames. + assert "x_embedder.weight" in converted + assert "context_embedder.weight" in converted + assert "proj_out.weight" in converted + + def test_converted_keys_are_all_real_transformer_params(self): + """Meta-device coverage: every converted key must exist in Flux2Transformer2DModel.""" + from diffusers import Flux2Transformer2DModel + + converted = convert_flux2_bfl_to_diffusers(keys_to_mock_state_dict(flux2_transformer_keys)) + + # The fixture keeps block 0 of each stack -> a single-layer model covers it. + with accelerate.init_empty_weights(): + model = Flux2Transformer2DModel(num_layers=1, num_single_layers=1) + params = set(model.state_dict().keys()) + + unmatched = sorted(k for k in converted if k not in params) + assert not unmatched, f"converted keys with no matching model parameter: {unmatched}" + + +class TestConvertFlux2Vae: + def test_full_bijective_coverage_against_arch(self): + """The full VAE fixture must convert to exactly the AutoencoderKLFlux2 parameter set.""" + from diffusers import AutoencoderKLFlux2 + + converted = convert_flux2_vae_bfl_to_diffusers(keys_to_mock_state_dict(flux2_vae_keys)) + + with accelerate.init_empty_weights(): + vae = AutoencoderKLFlux2(block_out_channels=(128, 256, 512, 512)) + params = set(vae.state_dict().keys()) + + unmatched = sorted(k for k in converted if k not in params) + missing = sorted(k for k in params if k not in converted) + assert not unmatched, f"converted keys with no matching VAE parameter: {unmatched}" + assert not missing, f"VAE parameters not covered by the converted checkpoint: {missing}" + + def test_up_block_order_is_reversed(self): + # BFL decoder.up.X maps to diffusers up_blocks.(3 - X). + sd = { + "decoder.up.0.block.0.norm1.weight": torch.empty(1), + "decoder.up.3.block.0.norm1.weight": torch.empty(1), + } + converted = convert_flux2_vae_bfl_to_diffusers(sd) + assert "decoder.up_blocks.3.resnets.0.norm1.weight" in converted + assert "decoder.up_blocks.0.resnets.0.norm1.weight" in converted + + def test_mid_attention_conv_weights_are_squeezed_to_linear(self): + # BFL stores mid attention as Conv2d [out, in, 1, 1]; diffusers uses Linear [out, in]. + sd = {"encoder.mid.attn_1.q.weight": torch.empty(8, 8, 1, 1)} + converted = convert_flux2_vae_bfl_to_diffusers(sd) + assert converted["encoder.mid_block.attentions.0.to_q.weight"].shape == (8, 8) + + +class TestSwapScaleShift: + def test_swaps_the_two_halves(self): + # First half = shift, second half = scale; diffusers wants them swapped. + weight = torch.cat([torch.zeros(2), torch.ones(2)]) # [shift=0, scale=1] + swapped = _flux2_swap_scale_shift(weight) + assert torch.allclose(swapped, torch.cat([torch.ones(2), torch.zeros(2)])) + + def test_leaves_malformed_tensor_untouched(self): + weight = torch.ones(3) # odd length -> cannot be split + assert torch.allclose(_flux2_swap_scale_shift(weight), weight) diff --git a/tests/backend/model_manager/load/test_qwen_image_state_dict_utils.py b/tests/backend/model_manager/load/test_qwen_image_state_dict_utils.py new file mode 100644 index 00000000000..425e9b6cc6b --- /dev/null +++ b/tests/backend/model_manager/load/test_qwen_image_state_dict_utils.py @@ -0,0 +1,169 @@ +"""Unit tests for the pure state-dict helpers in the Qwen-Image / Qwen-VL loader. + +These freeze the checkpoint key-surgery that the loaders perform before instantiating a model, +so a regression like the transformers-5.x one (where `_checkpoint_conversion_mapping` became +`{}` and the `visual.* -> model.visual.*` remap was silently skipped) fails here instead of at +the user's first load. +""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.qwen_image import ( + _build_qwen_image_transformer_config, + _dequantize_comfyui_fp8, + _remap_qwen_vl_checkpoint_keys, + _strip_comfyui_prefix, + _strip_quantization_metadata, +) +from tests.backend.model_manager.load.state_dicts.qwen_vl_encoder_comfyui_keys import ( + state_dict_keys as qwen_vl_keys, +) +from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict + +# Prefixes the Qwen2.5-VL architecture (transformers >=4.50) actually exposes. Frozen here on +# purpose: if a remap regresses, converted keys stop matching this set. +_VALID_QWEN_VL_PREFIXES = ("model.visual.", "model.language_model.", "lm_head") + + +class TestRemapQwenVlCheckpointKeys: + def test_every_key_maps_to_the_transformers_layout(self): + """Every legacy ComfyUI key must land under `model.visual.*` / `model.language_model.*`.""" + sd = keys_to_mock_state_dict(qwen_vl_keys) + # The loader strips fp8 metadata (`scaled_fp8` etc.) before remapping; mirror that order. + _strip_quantization_metadata(sd) + + remapped = _remap_qwen_vl_checkpoint_keys(sd) + + assert len(remapped) == len(sd) + for key in remapped: + assert key.startswith(_VALID_QWEN_VL_PREFIXES), f"key not remapped to a known layout: {key}" + # No key may survive in the legacy layout. + assert not any(k.startswith("visual.") for k in remapped) + assert not any(k.startswith(("model.layers.", "model.embed_tokens", "model.norm")) for k in remapped) + + def test_specific_keys_from_the_bug_report(self): + """The exact keys that failed to load in the original bug report are remapped.""" + sd = { + "visual.blocks.0.attn.qkv.weight": torch.empty(1), + "visual.patch_embed.proj.weight": torch.empty(1), + "model.layers.0.self_attn.q_proj.weight": torch.empty(1), + "model.embed_tokens.weight": torch.empty(1), + "lm_head.weight": torch.empty(1), + } + + remapped = _remap_qwen_vl_checkpoint_keys(sd) + + assert "model.visual.blocks.0.attn.qkv.weight" in remapped + assert "model.visual.patch_embed.proj.weight" in remapped + assert "model.language_model.layers.0.self_attn.q_proj.weight" in remapped + assert "model.language_model.embed_tokens.weight" in remapped + assert "lm_head.weight" in remapped # unchanged + + def test_idempotent_on_already_converted_layout(self): + """Re-running the remap on new-layout keys must not double-prefix them.""" + sd = keys_to_mock_state_dict(qwen_vl_keys) + + once = _remap_qwen_vl_checkpoint_keys(sd) + twice = _remap_qwen_vl_checkpoint_keys(once) + + assert set(once.keys()) == set(twice.keys()) + + def test_fallback_when_transformers_mapping_is_empty(self, monkeypatch): + """Even if transformers stops providing `_checkpoint_conversion_mapping`, the remap fires. + + transformers 5.x returns `{}` here; forcing that value pins the fallback that fixes the + original bug. + """ + from transformers import Qwen2_5_VLForConditionalGeneration + + monkeypatch.setattr(Qwen2_5_VLForConditionalGeneration, "_checkpoint_conversion_mapping", {}) + + remapped = _remap_qwen_vl_checkpoint_keys( + { + "visual.blocks.0.norm1.weight": torch.empty(1), + "model.layers.0.input_layernorm.weight": torch.empty(1), + } + ) + + assert "model.visual.blocks.0.norm1.weight" in remapped + assert "model.language_model.layers.0.input_layernorm.weight" in remapped + + +class TestStripQuantizationMetadata: + def test_drops_fp8_metadata_keeps_weights(self): + sd = keys_to_mock_state_dict(qwen_vl_keys) + # The captured checkpoint is fp8_scaled, so it really does ship this metadata. + assert any(k.endswith((".scale_weight", ".scale_input")) or k == "scaled_fp8" for k in sd) + n_weights_before = sum(1 for k in sd if k.endswith(".weight")) + + _strip_quantization_metadata(sd) + + assert not any( + k.endswith((".scale_weight", ".scale_input")) or "comfy_quant" in k or k == "scaled_fp8" for k in sd + ) + # Real weights are untouched. + assert sum(1 for k in sd if k.endswith(".weight")) == n_weights_before + + +class TestDequantizeComfyuiFp8: + def test_scalar_scale(self): + sd = { + "l.weight": torch.full((2, 2), 2.0), + "l.scale_weight": torch.tensor(3.0), + "l.scale_input": torch.tensor(9.0), # activation scale, must be ignored + } + + count = _dequantize_comfyui_fp8(sd, torch.float32) + + assert count == 1 + assert torch.allclose(sd["l.weight"], torch.full((2, 2), 6.0)) + + def test_block_wise_scale_is_broadcast(self): + # Per-block scale [2, 1] must be repeat_interleaved up to the weight shape [4, 2]. + sd = { + "l.weight": torch.ones(4, 2), + "l.weight_scale": torch.tensor([[10.0], [20.0]]), + } + + count = _dequantize_comfyui_fp8(sd, torch.float32) + + assert count == 1 + expected = torch.tensor([[10.0, 10.0], [10.0, 10.0], [20.0, 20.0], [20.0, 20.0]]) + assert torch.allclose(sd["l.weight"], expected) + + +class TestStripComfyuiPrefix: + def test_strips_diffusion_model_prefix(self): + sd = { + "model.diffusion_model.transformer_blocks.0.img_mod.1.weight": torch.empty(1), + "model.diffusion_model.img_in.weight": torch.empty(1), + } + out = _strip_comfyui_prefix(sd) + assert set(out.keys()) == {"transformer_blocks.0.img_mod.1.weight", "img_in.weight"} + + def test_no_prefix_is_a_noop(self): + sd = {"transformer_blocks.0.x": torch.empty(1)} + assert _strip_comfyui_prefix(sd) is sd + + +class TestBuildQwenImageTransformerConfig: + def test_infers_layer_count_and_dims_from_shapes(self): + # torch-order (logical) shapes, as the GGMLTensor.tensor_shape / safetensors path exposes. + sd = { + "img_in.weight": torch.empty(3072, 64), + "txt_in.weight": torch.empty(3072, 3584), + "transformer_blocks.0.img_mod.1.weight": torch.empty(1), + "transformer_blocks.1.img_mod.1.weight": torch.empty(1), + "transformer_blocks.5.img_mod.1.weight": torch.empty(1), + } + + cfg = _build_qwen_image_transformer_config(sd, is_edit=False) + + assert cfg["num_layers"] == 6 # max block index (5) + 1 + assert cfg["in_channels"] == 64 + assert cfg["num_attention_heads"] == 24 # 3072 // 128 + assert cfg["joint_attention_dim"] == 3584 + + def test_empty_state_dict_falls_back_to_defaults(self): + cfg = _build_qwen_image_transformer_config({}, is_edit=False) + assert cfg["num_layers"] == 60 diff --git a/tests/backend/model_manager/load/test_z_image_state_dict_utils.py b/tests/backend/model_manager/load/test_z_image_state_dict_utils.py new file mode 100644 index 00000000000..37e113d8570 --- /dev/null +++ b/tests/backend/model_manager/load/test_z_image_state_dict_utils.py @@ -0,0 +1,66 @@ +"""Unit tests for the Z-Image GGUF/ComfyUI -> diffusers state-dict converter.""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.z_image import _convert_z_image_gguf_to_diffusers +from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict +from tests.backend.model_manager.load.state_dicts.z_image_transformer_comfyui_keys import ( + state_dict_keys as z_image_keys, +) + + +class TestConvertZImageGgufToDiffusers: + def test_fused_qkv_split(self): + sd = keys_to_mock_state_dict(z_image_keys) + n_qkv = sum(1 for k in sd if k.endswith(".attention.qkv.weight")) + assert n_qkv > 0 + + out = _convert_z_image_gguf_to_diffusers(sd) + + # Each fused qkv weight becomes three separate projections. + assert sum(1 for k in out if k.endswith(".attention.to_q.weight")) == n_qkv + assert sum(1 for k in out if k.endswith(".attention.to_k.weight")) == n_qkv + assert sum(1 for k in out if k.endswith(".attention.to_v.weight")) == n_qkv + assert not any(".attention.qkv." in k for k in out) + + def test_key_renames(self): + out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys)) + # q_norm/k_norm -> norm_q/norm_k, attention.out -> attention.to_out.0 + assert any(k.endswith(".attention.norm_q.weight") for k in out) + assert any(k.endswith(".attention.norm_k.weight") for k in out) + assert any(k.endswith(".attention.to_out.0.weight") for k in out) + assert not any(".q_norm." in k or ".k_norm." in k for k in out) + assert not any(".attention.out." in k for k in out) + + def test_embedder_and_final_layer_renamed(self): + out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys)) + assert any(k.startswith("all_x_embedder.2-1.") for k in out) + assert any(k.startswith("all_final_layer.2-1.") for k in out) + assert not any(k.startswith("x_embedder.") or k.startswith("final_layer.") for k in out) + + def test_norm_final_is_dropped(self): + # The diffusers model uses a non-learnable final LayerNorm, so norm_final.* is skipped. + assert any(k.startswith("norm_final.") for k in z_image_keys) + out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys)) + assert not any(k.startswith("norm_final.") for k in out) + + def test_pad_tokens_are_2d_after_conversion(self): + # The diffusers model expects a leading batch dim on the pad tokens. The checkpoint + # already stores them 2D; GGUF ships them 1D (see the reshape test below). + out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys)) + for pad in ("x_pad_token", "cap_pad_token"): + assert out[pad].dim() == 2 + assert out[pad].shape[0] == 1 + + def test_1d_pad_token_gains_batch_dim(self): + # GGUF stores pad tokens as [dim]; they must be reshaped to [1, dim]. + out = _convert_z_image_gguf_to_diffusers({"x_pad_token": torch.arange(4.0)}) + assert out["x_pad_token"].shape == (1, 4) + + def test_qkv_split_preserves_values(self): + # A [6, 2] fused qkv splits into three [2, 2] chunks in order q, k, v. + qkv = torch.arange(12, dtype=torch.float32).reshape(6, 2) + out = _convert_z_image_gguf_to_diffusers({"blk.attention.qkv.weight": qkv}) + assert torch.allclose(out["blk.attention.to_q.weight"], qkv[0:2]) + assert torch.allclose(out["blk.attention.to_k.weight"], qkv[2:4]) + assert torch.allclose(out["blk.attention.to_v.weight"], qkv[4:6])