Upload sdxl_lora_elemental_tune.py

sdxl_lora_elemental_tune.py

@@ -7,178 +7,187 @@ 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___"
-          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
+    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
 
 
 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: ")
-
-  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: ")
+    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]:
+                                if target_key.endswith(".alpha"):
+                                  final_weight = weight * multiplier
+                                  if not remove_zero_weight_keys or final_weight != 0.0:
+                                    adjusted_tensors[target_key] = lora_tensors[target_key]
+                                else:
+                                  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)
+
+        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.endswith(".alpha"):
+                                      final_weight = weight * multiplier
+                                      if not remove_zero_weight_keys or final_weight != 0.0:
+                                        adjusted_tensors[target_key] = lora_tensors[target_key]
+                                    else:
+                                      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: {e}")
 
 
 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: ")
+    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}")
 
 
 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: ")
+    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}")
 
 if __name__ == "__main__":
-  main()
+    main()