|
|
import safetensors.torch
|
|
|
from safetensors import safe_open
|
|
|
import torch
|
|
|
|
|
|
def patch_final_layer_adaLN(state_dict, prefix="lora_unet_final_layer", verbose=True):
|
|
|
"""
|
|
|
Add dummy adaLN weights if missing, using final_layer_linear shapes as reference.
|
|
|
Args:
|
|
|
state_dict (dict): keys -> tensors
|
|
|
prefix (str): base name for final_layer keys
|
|
|
verbose (bool): print debug info
|
|
|
Returns:
|
|
|
dict: patched state_dict
|
|
|
"""
|
|
|
final_layer_linear_down = None
|
|
|
final_layer_linear_up = None
|
|
|
|
|
|
adaLN_down_key = f"{prefix}_adaLN_modulation_1.lora_down.weight"
|
|
|
adaLN_up_key = f"{prefix}_adaLN_modulation_1.lora_up.weight"
|
|
|
linear_down_key = f"{prefix}_linear.lora_down.weight"
|
|
|
linear_up_key = f"{prefix}_linear.lora_up.weight"
|
|
|
|
|
|
if verbose:
|
|
|
print(f"\n🔍 Checking for final_layer keys with prefix: '{prefix}'")
|
|
|
print(f" Linear down: {linear_down_key}")
|
|
|
print(f" Linear up: {linear_up_key}")
|
|
|
|
|
|
if linear_down_key in state_dict:
|
|
|
final_layer_linear_down = state_dict[linear_down_key]
|
|
|
if linear_up_key in state_dict:
|
|
|
final_layer_linear_up = state_dict[linear_up_key]
|
|
|
|
|
|
has_adaLN = adaLN_down_key in state_dict and adaLN_up_key in state_dict
|
|
|
has_linear = final_layer_linear_down is not None and final_layer_linear_up is not None
|
|
|
|
|
|
if verbose:
|
|
|
print(f" ✅ Has final_layer.linear: {has_linear}")
|
|
|
print(f" ✅ Has final_layer.adaLN_modulation_1: {has_adaLN}")
|
|
|
|
|
|
if has_linear and not has_adaLN:
|
|
|
dummy_down = torch.zeros_like(final_layer_linear_down)
|
|
|
dummy_up = torch.zeros_like(final_layer_linear_up)
|
|
|
state_dict[adaLN_down_key] = dummy_down
|
|
|
state_dict[adaLN_up_key] = dummy_up
|
|
|
|
|
|
if verbose:
|
|
|
print(f"✅ Added dummy adaLN weights:")
|
|
|
print(f" {adaLN_down_key} (shape: {dummy_down.shape})")
|
|
|
print(f" {adaLN_up_key} (shape: {dummy_up.shape})")
|
|
|
else:
|
|
|
if verbose:
|
|
|
print("✅ No patch needed — adaLN weights already present or no final_layer.linear found.")
|
|
|
|
|
|
return state_dict
|
|
|
|
|
|
|
|
|
def main():
|
|
|
print("🔄 Universal final_layer.adaLN LoRA patcher (.safetensors)")
|
|
|
input_path = input("Enter path to input LoRA .safetensors file: ").strip()
|
|
|
output_path = input("Enter path to save patched LoRA .safetensors file: ").strip()
|
|
|
|
|
|
|
|
|
state_dict = {}
|
|
|
with safe_open(input_path, framework="pt", device="cpu") as f:
|
|
|
for k in f.keys():
|
|
|
state_dict[k] = f.get_tensor(k)
|
|
|
|
|
|
print(f"\n✅ Loaded {len(state_dict)} tensors from: {input_path}")
|
|
|
|
|
|
|
|
|
final_keys = [k for k in state_dict if "final_layer" in k]
|
|
|
if final_keys:
|
|
|
print("\n🔑 Found these final_layer-related keys:")
|
|
|
for k in final_keys:
|
|
|
print(f" {k}")
|
|
|
else:
|
|
|
print("\n⚠️ No keys with 'final_layer' found — will try patch anyway.")
|
|
|
|
|
|
|
|
|
prefixes = [
|
|
|
"lora_unet_final_layer",
|
|
|
"final_layer",
|
|
|
"base_model.model.final_layer"
|
|
|
]
|
|
|
patched = False
|
|
|
|
|
|
for prefix in prefixes:
|
|
|
before = len(state_dict)
|
|
|
state_dict = patch_final_layer_adaLN(state_dict, prefix=prefix)
|
|
|
after = len(state_dict)
|
|
|
if after > before:
|
|
|
patched = True
|
|
|
break
|
|
|
|
|
|
if not patched:
|
|
|
print("\nℹ️ No patch applied — either adaLN already exists or no final_layer.linear found.")
|
|
|
|
|
|
|
|
|
safetensors.torch.save_file(state_dict, output_path)
|
|
|
print(f"\n✅ Patched file saved to: {output_path}")
|
|
|
print(f" Total tensors now: {len(state_dict)}")
|
|
|
|
|
|
|
|
|
print("\n🔍 Verifying patched keys:")
|
|
|
with safe_open(output_path, framework="pt", device="cpu") as f:
|
|
|
keys = list(f.keys())
|
|
|
for k in keys:
|
|
|
if "final_layer" in k:
|
|
|
print(f" {k}")
|
|
|
|
|
|
has_adaLN_after = any("adaLN_modulation_1" in k for k in keys)
|
|
|
print(f"✅ Contains adaLN after patch: {has_adaLN_after}")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main()
|
|
|
|
