sdxl_lora_elemental_tune.pyの強度適用する処理修正。

sdxl_lora_elemental_tune.py

@@ -4,175 +4,181 @@ from safetensors.torch import load_file, save_file
 import toml
 import re
 from safetensors import safe_open
+import math
 
 def parse_key(key):
-    match = re.match(r"lora_unet_(input|output|up|down)_blocks_(\d+(?:_\d+)?)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
-    if match:
-        return "unet", match.group(1) + "_blocks", match.group(2), match.group(3)
+  match = re.match(r"lora_unet_(input|output|up|down)_blocks_(\d+(?:_\d+)?)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
+  if match:
+    return "unet", match.group(1) + "_blocks", match.group(2), match.group(3)
 
-    match = re.match(r"lora_unet_(mid_block)_(resnets|attentions)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
-    if match:
-        return "unet", match.group(1), f"{match.group(2)}_{match.group(3)}", match.group(4)
+  match = re.match(r"lora_unet_(mid_block)_(resnets|attentions)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
+  if match:
+    return "unet", match.group(1), f"{match.group(2)}_{match.group(3)}", match.group(4)
 
-    match = re.match(r"lora_unet_(middle_block)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
-    if match:
-        return "unet", match.group(1), match.group(2), match.group(3)
+  match = re.match(r"lora_unet_(middle_block)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
+  if match:
+    return "unet", match.group(1), match.group(2), match.group(3)
 
-    match = re.match(r"lora_te\d+_text_model_encoder_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
-    if match:
-        return "text_encoder", "encoder_layers", match.group(1).split("_")[0], "_".join(match.group(1).split("_")[1:])
+  match = re.match(r"lora_te\d+_text_model_encoder_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
+  if match:
+    return "text_encoder", "encoder_layers", match.group(1).split("_")[0], "_".join(match.group(1).split("_")[1:])
 
-    return None, None, None, None
+  return None, None, None, None
 
 def extract_lora_hierarchy(lora_tensors, mode="extract"):
-    lora_hierarchy = {}
-    lora_key_groups = {"unet": {}, "text_encoder": {}} if mode == "adjust" else None
-
-    for key in lora_tensors:
-        if key.startswith("lora_unet_"):
-            model_type, block_type, block_num, layer_key = parse_key(key)
-
-            if model_type and block_type and layer_key:
-                parts = layer_key.split("_")
-                if "transformer_blocks" in layer_key:
-                    grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
-                elif "attentions" in layer_key:
-                    grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
-                elif "resnets" in layer_key:
-                    grouped_key = "_".join(parts[:3])
-                else:
-                    grouped_key = layer_key
-
-                if model_type not in lora_hierarchy:
-                    lora_hierarchy[model_type] = {}
-                if block_type not in lora_hierarchy[model_type]:
-                    lora_hierarchy[model_type][block_type] = {}
-                if block_num not in lora_hierarchy[model_type][block_type]:
-                    lora_hierarchy[model_type][block_type][block_num] = {}
-                lora_hierarchy[model_type][block_type][block_num][grouped_key] = 1.0
-
-                if mode == "adjust":
-                    group_key = f"..unet_{block_type}_{block_num}_{grouped_key}"
-                    if group_key not in lora_key_groups["unet"]:
-                        lora_key_groups["unet"][group_key] = []
-                    lora_key_groups["unet"][group_key].append(key)
-
-        elif key.startswith("lora_te"):
-            match = re.match(r"(lora_te\d+)_text_model_encoder_layers_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
-            if match:
-                model_section = match.group(1)
-                block_type = "encoder"
-                block_num = match.group(2)
-                layer_key = match.group(3)
-
-                grouped_key = f"layers_{block_num}__{layer_key}"
-
-                if model_section not in lora_hierarchy:
-                    lora_hierarchy[model_section] = {}
-                if block_type not in lora_hierarchy[model_section]:
-                    lora_hierarchy[model_section][block_type] = {}
-                lora_hierarchy[model_section][block_type][grouped_key] = 1.0
-
-                if mode == "adjust":
-                    group_key = f"..{model_section}_{block_num}_{layer_key}"
-                    lora_key_groups["text_encoder"][group_key] = [key]
-
-    return lora_hierarchy if mode == "extract" else lora_key_groups
+  lora_hierarchy = {}
+  lora_key_groups = {"unet": {}, "text_encoder": {}} if mode == "adjust" else None
+
+  for key in lora_tensors:
+    if key.startswith("lora_unet_"):
+      model_type, block_type, block_num, layer_key = parse_key(key)
+
+      if model_type and block_type and layer_key:
+        parts = layer_key.split("_")
+        if "transformer_blocks" in layer_key:
+          grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
+        elif "attentions" in layer_key:
+          grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
+        elif "resnets" in layer_key:
+          grouped_key = "_".join(parts[:3])
+        else:
+          grouped_key = layer_key
+
+        if model_type not in lora_hierarchy:
+          lora_hierarchy[model_type] = {}
+        if block_type not in lora_hierarchy[model_type]:
+          lora_hierarchy[model_type][block_type] = {}
+        if block_num not in lora_hierarchy[model_type][block_type]:
+          lora_hierarchy[model_type][block_type][block_num] = {}
+        lora_hierarchy[model_type][block_type][block_num][grouped_key] = 1.0
+
+        if mode == "adjust":
+          group_key = f"..unet___"
+          if group_key not in lora_key_groups["unet"]:
+            lora_key_groups["unet"][group_key] = []
+          lora_key_groups["unet"][group_key].append(key)
+
+    elif key.startswith("lora_te"):
+      match = re.match(r"(lora_te\d+)_text_model_encoder_layers_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
+      if match:
+        model_section = match.group(1)
+        block_type = "encoder"
+        block_num = match.group(2)
+        layer_key = match.group(3)
+
+        grouped_key = f"layers___"
+
+        if model_section not in lora_hierarchy:
+          lora_hierarchy[model_section] = {}
+        if block_type not in lora_hierarchy[model_section]:
+          lora_hierarchy[model_section][block_type] = {}
+        lora_hierarchy[model_section][block_type][grouped_key] = 1.0
+
+        if mode == "adjust":
+          group_key = f"..__"
+          lora_key_groups["text_encoder"][group_key] = [key]
+
+  return lora_hierarchy if mode == "extract" else lora_key_groups
 
 
 def adjust_lora_weights(lora_path, toml_path, output_path, multiplier=1.0, remove_zero_weight_keys=True):
-    try:
-        lora_tensors = load_file(lora_path)
-        with safe_open(lora_path, framework="pt") as f:
-            metadata = f.metadata()
-    except Exception as e:
-        raise Exception(f"Error loading LoRA model: {e}")
-
-    try:
-        with open(toml_path, "r") as f:
-            lora_config = toml.load(f)
-    except Exception as e:
-        raise Exception(f"Error loading TOML file: {e}")
-
-
-    lora_key_groups = extract_lora_hierarchy(lora_tensors, mode="adjust")
-    adjusted_tensors = {}
-
-    for model_section, model_config in lora_config.items():
-        if model_section.startswith("lora_te"):
-            for block_type, layers in model_config.items():
-                for layer_key, weight in layers.items():
-                    block_num, layer_name = layer_key.replace("layers_", "").split("__")
-                    group_key = f"..{model_section}_{block_num}_{layer_name}"
-                    if group_key in lora_key_groups["text_encoder"]:
-                        final_weight = weight * multiplier
-                        if not remove_zero_weight_keys or final_weight != 0.0:
-                            for target_key in lora_key_groups["text_encoder"][group_key]:
-                                adjusted_tensors[target_key] = lora_tensors[target_key] * final_weight
-
-        else:  # unet
-            for block_type, block_nums in model_config.items():
-                for block_num, layer_keys in block_nums.items():
-                    for grouped_key, weight in layer_keys.items():
-                        group_key = f"..unet_{block_type}_{block_num}_{grouped_key}"
-                        if group_key in lora_key_groups["unet"]:
-                            final_weight = weight * multiplier
-                            if not remove_zero_weight_keys or final_weight != 0.0:
-                                for target_key in lora_key_groups["unet"][group_key]:
-                                    if target_key in lora_tensors:
-                                        adjusted_tensors[target_key] = lora_tensors[target_key] * final_weight
-
-
-    try:
-        save_file(adjusted_tensors, output_path, metadata)
-    except Exception as e:
-        raise Exception(f"Error saving adjusted model: {e}")
+  try:
+    lora_tensors = load_file(lora_path)
+    with safe_open(lora_path, framework="pt") as f:
+      metadata = f.metadata()
+  except Exception as e:
+    raise Exception(f"Error loading LoRA model: ")
+
+  try:
+    with open(toml_path, "r") as f:
+      lora_config = toml.load(f)
+  except Exception as e:
+    raise Exception(f"Error loading TOML file: ")
+
+
+  lora_key_groups = extract_lora_hierarchy(lora_tensors, mode="adjust")
+  adjusted_tensors = {}
+
+  for model_section, model_config in lora_config.items():
+    if model_section.startswith("lora_te"):
+      for block_type, layers in model_config.items():
+        for layer_key, weight in layers.items():
+          block_num, layer_name = layer_key.replace("layers_", "").split("__")
+          group_key = f"..__"
+          if group_key in lora_key_groups["text_encoder"]:
+            final_weight = weight * multiplier
+            if not remove_zero_weight_keys or final_weight != 0.0:
+              for target_key in lora_key_groups["text_encoder"][group_key]:
+                adjusted_tensors[target_key] = lora_tensors[target_key] * final_weight
+
+    else:  # unet
+      for block_type, block_nums in model_config.items():
+        for block_num, layer_keys in block_nums.items():
+          for grouped_key, weight in layer_keys.items():
+            group_key = f"..unet___"
+            if group_key in lora_key_groups["unet"]:
+              for target_key in lora_key_groups["unet"][group_key]:
+                if target_key in lora_tensors:
+                  if target_key.endswith(".alpha"):
+                    # .alpha は変更しない
+                    adjusted_tensors[target_key] = lora_tensors[target_key]
+                  else:
+                    # lora_down.weight, lora_up.weight は平方根を適用
+                    final_weight = weight * multiplier
+                    if not remove_zero_weight_keys or final_weight != 0.0:
+                      adjusted_tensors[target_key] = lora_tensors[target_key] * math.sqrt(final_weight)
+
+
+  try:
+    save_file(adjusted_tensors, output_path, metadata)
+  except Exception as e:
+    raise Exception(f"Error saving adjusted model: ")
 
 
 def write_toml(lora_hierarchy, output_path):
-    try:
-        with open(output_path, "w") as f:
-            toml.dump(lora_hierarchy, f)
-    except Exception as e:
-        raise Exception(f"Error writing TOML file: {e}")
+  try:
+    with open(output_path, "w") as f:
+      toml.dump(lora_hierarchy, f)
+  except Exception as e:
+    raise Exception(f"Error writing TOML file: ")
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Extract or adjust LoRA weights based on a TOML config.")
-    subparsers = parser.add_subparsers(dest="mode", help="Choose mode: 'extract' or 'adjust'")
-
-    # Extract mode
-    parser_extract = subparsers.add_parser("extract", help="Extract LoRA hierarchy to a TOML file")
-    parser_extract.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
-    parser_extract.add_argument("--output_path", required=True, help="Path to the output TOML file")
-
-    # Adjust mode
-    parser_adjust = subparsers.add_parser("adjust", help="Adjust LoRA weights based on a TOML config.")
-    parser_adjust.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
-    parser_adjust.add_argument("--toml_path", required=True, help="Path to the TOML config file")
-    parser_adjust.add_argument("--output_path", required=True, help="Path to the output safetensors file")
-    parser_adjust.add_argument("--multiplier", type=float, default=1.0, help="Global multiplier for the LoRA weights")
-    parser_adjust.add_argument("--remove_zero_weight_keys", action="store_true",
-                            help="Remove keys with resulting weight of 0. Useful for reducing file size.")
-
-    args = parser.parse_args()
-
-    try:
-        if args.mode == "extract":
-            lora_tensors = load_file(args.lora_path)
-            lora_hierarchy = extract_lora_hierarchy(lora_tensors) 
-            write_toml(lora_hierarchy, args.output_path)
-            print(f"Successfully extracted LoRA hierarchy to {args.output_path}")
-
-        elif args.mode == "adjust":
-            adjust_lora_weights(args.lora_path, args.toml_path, args.output_path, args.multiplier, args.remove_zero_weight_keys)
-            print(f"Successfully adjusted LoRA weights and saved to {args.output_path}")
-
-        else:
-            parser.print_help()
-
-    except Exception as e:
-        print(f"An error occurred: {e}")
+  parser = argparse.ArgumentParser(description="Extract or adjust LoRA weights based on a TOML config.")
+  subparsers = parser.add_subparsers(dest="mode", help="Choose mode: 'extract' or 'adjust'")
+
+  # Extract mode
+  parser_extract = subparsers.add_parser("extract", help="Extract LoRA hierarchy to a TOML file")
+  parser_extract.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
+  parser_extract.add_argument("--output_path", required=True, help="Path to the output TOML file")
+
+  # Adjust mode
+  parser_adjust = subparsers.add_parser("adjust", help="Adjust LoRA weights based on a TOML config.")
+  parser_adjust.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
+  parser_adjust.add_argument("--toml_path", required=True, help="Path to the TOML config file")
+  parser_adjust.add_argument("--output_path", required=True, help="Path to the output safetensors file")
+  parser_adjust.add_argument("--multiplier", type=float, default=1.0, help="Global multiplier for the LoRA weights")
+  parser_adjust.add_argument("--remove_zero_weight_keys", action="store_true",
+                              help="Remove keys with resulting weight of 0. Useful for reducing file size.")
+
+  args = parser.parse_args()
+
+  try:
+    if args.mode == "extract":
+      lora_tensors = load_file(args.lora_path)
+      lora_hierarchy = extract_lora_hierarchy(lora_tensors) 
+      write_toml(lora_hierarchy, args.output_path)
+      print(f"Successfully extracted LoRA hierarchy to {args.output_path}")
+
+    elif args.mode == "adjust":
+      adjust_lora_weights(args.lora_path, args.toml_path, args.output_path, args.multiplier, args.remove_zero_weight_keys)
+      print(f"Successfully adjusted LoRA weights and saved to {args.output_path}")
+
+    else:
+      parser.print_help()
+
+  except Exception as e:
+    print(f"An error occurred: ")
 
 if __name__ == "__main__":
-    main()
+  main()