mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-04-23 03:00:31 -04:00
Add unit test for lora_model_from_flux_diffusers_state_dict(...).
This commit is contained in:
Ryan Dick
@@ -8,211 +8,8 @@ from invokeai.backend.lora.layers.lora_layer import LoRALayer
|
||||
from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
|
||||
from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
|
||||
|
||||
# def convert_flux_transformer_checkpoint_to_diffusers(
|
||||
# original_state_dict, num_layers, num_single_layers, inner_dim, mlp_ratio=4.0
|
||||
# ):
|
||||
# converted_state_dict = {}
|
||||
|
||||
# ## time_text_embed.timestep_embedder <- time_in
|
||||
# converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop(
|
||||
# "time_in.in_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop(
|
||||
# "time_in.in_layer.bias"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop(
|
||||
# "time_in.out_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop(
|
||||
# "time_in.out_layer.bias"
|
||||
# )
|
||||
|
||||
# ## time_text_embed.text_embedder <- vector_in
|
||||
# converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop(
|
||||
# "vector_in.in_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop(
|
||||
# "vector_in.in_layer.bias"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop(
|
||||
# "vector_in.out_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop(
|
||||
# "vector_in.out_layer.bias"
|
||||
# )
|
||||
|
||||
# # guidance
|
||||
# has_guidance = any("guidance" in k for k in original_state_dict)
|
||||
# if has_guidance:
|
||||
# converted_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = original_state_dict.pop(
|
||||
# "guidance_in.in_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = original_state_dict.pop(
|
||||
# "guidance_in.in_layer.bias"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = original_state_dict.pop(
|
||||
# "guidance_in.out_layer.weight"
|
||||
# )
|
||||
# converted_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = original_state_dict.pop(
|
||||
# "guidance_in.out_layer.bias"
|
||||
# )
|
||||
|
||||
# # context_embedder
|
||||
# converted_state_dict["context_embedder.weight"] = original_state_dict.pop("txt_in.weight")
|
||||
# converted_state_dict["context_embedder.bias"] = original_state_dict.pop("txt_in.bias")
|
||||
|
||||
# # x_embedder
|
||||
# converted_state_dict["x_embedder.weight"] = original_state_dict.pop("img_in.weight")
|
||||
# converted_state_dict["x_embedder.bias"] = original_state_dict.pop("img_in.bias")
|
||||
|
||||
# # double transformer blocks
|
||||
# for i in range(num_layers):
|
||||
# block_prefix = f"transformer_blocks.{i}."
|
||||
# # norms.
|
||||
# ## norm1
|
||||
# converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mod.lin.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mod.lin.bias"
|
||||
# )
|
||||
# ## norm1_context
|
||||
# converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mod.lin.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mod.lin.bias"
|
||||
# )
|
||||
# # Q, K, V
|
||||
# sample_q, sample_k, sample_v = torch.chunk(
|
||||
# original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0
|
||||
# )
|
||||
# context_q, context_k, context_v = torch.chunk(
|
||||
# original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0
|
||||
# )
|
||||
# sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk(
|
||||
# original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0
|
||||
# )
|
||||
# context_q_bias, context_k_bias, context_v_bias = torch.chunk(
|
||||
# original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v])
|
||||
# converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias])
|
||||
# # qk_norm
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_attn.norm.query_norm.scale"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_attn.norm.key_norm.scale"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_attn.norm.query_norm.scale"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_attn.norm.key_norm.scale"
|
||||
# )
|
||||
# # ff img_mlp
|
||||
# converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mlp.0.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mlp.0.bias"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mlp.2.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_mlp.2.bias"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mlp.0.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mlp.0.bias"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mlp.2.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_mlp.2.bias"
|
||||
# )
|
||||
# # output projections.
|
||||
# converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_attn.proj.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.img_attn.proj.bias"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_attn.proj.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop(
|
||||
# f"double_blocks.{i}.txt_attn.proj.bias"
|
||||
# )
|
||||
|
||||
# # single transfomer blocks
|
||||
# for i in range(num_single_layers):
|
||||
# block_prefix = f"single_transformer_blocks.{i}."
|
||||
# # norm.linear <- single_blocks.0.modulation.lin
|
||||
# converted_state_dict[f"{block_prefix}norm.linear.weight"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.modulation.lin.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}norm.linear.bias"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.modulation.lin.bias"
|
||||
# )
|
||||
# # Q, K, V, mlp
|
||||
# mlp_hidden_dim = int(inner_dim * mlp_ratio)
|
||||
# split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim)
|
||||
# q, k, v, mlp = torch.split(original_state_dict.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0)
|
||||
# q_bias, k_bias, v_bias, mlp_bias = torch.split(
|
||||
# original_state_dict.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v])
|
||||
# converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias])
|
||||
# converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp])
|
||||
# converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias])
|
||||
# # qk norm
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.norm.query_norm.scale"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.norm.key_norm.scale"
|
||||
# )
|
||||
# # output projections.
|
||||
# converted_state_dict[f"{block_prefix}proj_out.weight"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.linear2.weight"
|
||||
# )
|
||||
# converted_state_dict[f"{block_prefix}proj_out.bias"] = original_state_dict.pop(
|
||||
# f"single_blocks.{i}.linear2.bias"
|
||||
# )
|
||||
|
||||
# converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight")
|
||||
# converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias")
|
||||
# converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(
|
||||
# original_state_dict.pop("final_layer.adaLN_modulation.1.weight")
|
||||
# )
|
||||
# converted_state_dict["norm_out.linear.bias"] = swap_scale_shift(
|
||||
# original_state_dict.pop("final_layer.adaLN_modulation.1.bias")
|
||||
# )
|
||||
|
||||
# return converted_state_dict
|
||||
|
||||
|
||||
# TODO(ryand): What alpha should we use? 1.0? Rank of the matrix?
|
||||
# TODO(ryand): What alpha should we use? 1.0? Rank of the LoRA?
|
||||
def lora_model_from_flux_diffusers_state_dict(state_dict: Dict[str, torch.Tensor], alpha: float = 1.0) -> LoRAModelRaw: # pyright: ignore[reportRedeclaration] (state_dict is intentionally re-declared)
|
||||
"""Loads a state dict in the Diffusers FLUX LoRA format into a LoRAModelRaw object.
|
||||
|
||||
@@ -222,6 +19,9 @@ def lora_model_from_flux_diffusers_state_dict(state_dict: Dict[str, torch.Tensor
|
||||
# Group keys by layer.
|
||||
grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_by_layer(state_dict)
|
||||
|
||||
# Remove the "transformer." prefix from all keys.
|
||||
grouped_state_dict = {k.replace("transformer.", ""): v for k, v in grouped_state_dict.items()}
|
||||
|
||||
# Constants for FLUX.1
|
||||
num_double_layers = 19
|
||||
num_single_layers = 38
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
import torch
|
||||
|
||||
from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils import (
|
||||
lora_model_from_flux_diffusers_state_dict,
|
||||
)
|
||||
from tests.backend.lora.conversions.lora_state_dicts.flux_lora_diffusers_format import state_dict_keys
|
||||
|
||||
|
||||
def test_lora_model_from_flux_diffusers_state_dict():
|
||||
"""Test that lora_model_from_flux_diffusers_state_dict() can load a state dict in the Diffusers FLUX LoRA format."""
|
||||
# Construct a state dict that is in the Diffusers FLUX LoRA format.
|
||||
state_dict: dict[str, torch.Tensor] = {}
|
||||
for k in state_dict_keys:
|
||||
state_dict[k] = torch.empty(1)
|
||||
|
||||
# Load the state dict into a LoRAModelRaw object.
|
||||
model = lora_model_from_flux_diffusers_state_dict(state_dict)
|
||||
|
||||
# Check that the model has the correct number of LoRA layers.
|
||||
expected_lora_layers: set[str] = set()
|
||||
for k in state_dict_keys:
|
||||
k = k.replace("lora_A.weight", "")
|
||||
k = k.replace("lora_B.weight", "")
|
||||
expected_lora_layers.add(k)
|
||||
# Drop the K/V/proj_mlp weights because these are all concatenated into a single layer in the BFL format (we keep
|
||||
# the Q weights so that we count these layers once).
|
||||
concatenated_weights = ["to_k", "to_v", "proj_mlp", "add_k_proj", "add_v_proj"]
|
||||
expected_lora_layers = {k for k in expected_lora_layers if not any(w in k for w in concatenated_weights)}
|
||||
assert len(model.layers) == len(expected_lora_layers)
|
||||
Reference in New Issue
Block a user