Skip to content

[WIP] [LoRA] support omi hidream lora. #11660

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 52 additions & 6 deletions src/diffusers/loaders/lora_conversion_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1789,12 +1789,58 @@ def get_alpha_scales(down_weight, key):
return converted_state_dict


def _convert_non_diffusers_hidream_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"):
if not all(k.startswith(non_diffusers_prefix) for k in state_dict):
raise ValueError("Invalid LoRA state dict for HiDream.")
converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()}
converted_state_dict = {f"transformer.{k}": v for k, v in converted_state_dict.items()}
return converted_state_dict
def _convert_non_diffusers_hidream_lora_to_diffusers(state_dict):
non_diffusers_prefix = "diffusion_model"
is_kohya = all(k.startswith(f"{non_diffusers_prefix}.") for k in state_dict)

def _convert_kohya(state_dict):
converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()}
converted_state_dict = {f"transformer.{k}": v for k, v in converted_state_dict.items()}
return converted_state_dict

if is_kohya:
return _convert_kohya(state_dict)

else:
assert any(k.startswith(("clip_g.", "clip_l.", "t5.", "llama.", "transformer.")) for k in state_dict)
converted_state_dict = {}
component = "transformer"
compoent_sd = {k: v for k, v in state_dict.items() if k.startswith(f"{component}.")}

def _convert_omi(key, state_dict, component):
down_key = f"{key}.lora_down.weight"
down_weight = state_dict.pop(down_key)
lora_rank = down_weight.shape[0]

up_weight_key = f"{key}.lora_up.weight"
up_weight = state_dict.pop(up_weight_key)

alpha_key = f"{key}.alpha"
alpha = state_dict.pop(alpha_key)

# scale weight by alpha and dim
scale = alpha / lora_rank
# calculate scale_down and scale_up
scale_down = scale
scale_up = 1.0
while scale_down * 2 < scale_up:
scale_down *= 2
scale_up /= 2
down_weight = down_weight * scale_down
up_weight = up_weight * scale_up

diffusers_down_key = f"{key}.lora_A.weight"
converted_state_dict[f"{component}.{diffusers_down_key}"] = down_weight
converted_state_dict[f"{component}.{diffusers_down_key.replace('.lora_A.', '.lora_B.')}"] = up_weight

all_unique_keys = {
k.replace(".lora_down.weight", "").replace(".lora_up.weight", "").replace(".alpha", "")
for k in compoent_sd
}
for k in all_unique_keys:
_convert_omi(k, compoent_sd, component=component)

return converted_state_dict


def _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"):
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/loaders/lora_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -5489,7 +5489,9 @@ def lora_state_dict(
logger.warning(warn_msg)
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}

is_non_diffusers_format = any("diffusion_model" in k for k in state_dict)
kohya_format = any("diffusion_model" in k for k in state_dict)
is_omi_format = any(k.startswith(("clip_g.", "clip_l.", "t5.", "llama.", "transformer.")) for k in state_dict)
is_non_diffusers_format = kohya_format or is_omi_format
if is_non_diffusers_format:
state_dict = _convert_non_diffusers_hidream_lora_to_diffusers(state_dict)

Expand Down
Loading