| 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()
|
|
|
