60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
import torch
|
|
from comfy.model_detection import count_blocks
|
|
from comfy.utils import flux_to_diffusers
|
|
|
|
def convert_diffusers_flux_lora(state_dict, output_prefix=""):
|
|
out_sd = {}
|
|
|
|
state_dict = {key.replace('transformer.', ''): value for key, value in state_dict.items()}
|
|
|
|
new_sd = {}
|
|
for lora_a in state_dict.keys():
|
|
if "lora_A" not in lora_a:
|
|
continue
|
|
lora_b = lora_a.replace("lora_A", "lora_B")
|
|
key = lora_a.replace("lora_A.", "")
|
|
assert "lora_A" in lora_a and "lora_B" in lora_b, f"Invalid LoRA checkpoint. {lora_a} {lora_b}"
|
|
|
|
new_sd[key] = state_dict[lora_b] @ state_dict[lora_a]
|
|
|
|
state_dict = new_sd
|
|
|
|
depth = count_blocks(state_dict, 'transformer_blocks.{}.')
|
|
depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.')
|
|
hidden_size = state_dict["x_embedder.weight"].shape[0]
|
|
sd_map = flux_to_diffusers({"depth": depth, "depth_single_blocks": depth_single_blocks, "hidden_size": hidden_size}, output_prefix=output_prefix)
|
|
|
|
for k in sd_map:
|
|
weight = state_dict.get(k, None)
|
|
if weight is not None:
|
|
t = sd_map[k]
|
|
|
|
if not isinstance(t, str):
|
|
if len(t) > 2:
|
|
fun = t[2]
|
|
else:
|
|
fun = lambda a: a
|
|
offset = t[1]
|
|
if offset is not None:
|
|
old_weight = out_sd.get(t[0], None)
|
|
if old_weight is None:
|
|
old_weight = torch.empty_like(weight)
|
|
if old_weight.shape[offset[0]] < offset[1] + offset[2]:
|
|
exp = list(weight.shape)
|
|
exp[offset[0]] = offset[1] + offset[2]
|
|
new = torch.empty(exp, device=weight.device, dtype=weight.dtype)
|
|
new[:old_weight.shape[0]] = old_weight
|
|
old_weight = new
|
|
|
|
w = old_weight.narrow(offset[0], offset[1], offset[2])
|
|
else:
|
|
old_weight = weight
|
|
w = weight
|
|
w[:] = fun(weight)
|
|
t = t[0]
|
|
out_sd[t] = old_weight
|
|
else:
|
|
out_sd[t] = weight
|
|
state_dict.pop(k)
|
|
|
|
return out_sd |