From cbe2a9da45fd99068faf1511b7f6bf1e641dd6d4 Mon Sep 17 00:00:00 2001
From: Kohya S <52813779+kohya-ss@users.noreply.github.com>
Date: Tue, 16 Sep 2025 21:48:47 +0900
Subject: [PATCH] feat: add conversion script for LoRA models to ComfyUI format
 with reverse option

---
 docs/hunyuan_image_train_network.md           | 20 ++++++-
 .../convert_hunyuan_image_lora_to_comfy.py    | 54 +++++++++++++------
 2 files changed, 55 insertions(+), 19 deletions(-)

diff --git a/docs/hunyuan_image_train_network.md b/docs/hunyuan_image_train_network.md
index c4c93d8d..3d49fbdf 100644
--- a/docs/hunyuan_image_train_network.md
+++ b/docs/hunyuan_image_train_network.md
@@ -461,12 +461,28 @@ python hunyuan_image_minimal_inference.py \
 
 ## 9. Related Tools / 関連ツール
 
-- **`hunyuan_image_minimal_inference.py`**: Simple inference script for generating images with trained LoRA models.
+### `networks/convert_hunyuan_image_lora_to_comfy.py`
+
+A script to convert LoRA models to ComfyUI-compatible format. The formats differ slightly, so conversion is necessary. You can convert from the sd-scripts format to ComfyUI format with:
+
+```bash
+python networks/convert_hunyuan_image_lora_to_comfy.py path/to/source.safetensors path/to/destination.safetensors
+```
+
+Using the `--reverse` option allows conversion in the opposite direction (ComfyUI format to sd-scripts format). However, reverse conversion is only possible for LoRAs converted by this script. LoRAs created with other training tools cannot be converted.
 
 <details>
 <summary>日本語</summary>
 
-- **`hunyuan_image_minimal_inference.py`**: 学習した LoRA モデルを適用して画像を生成するシンプルな推論スクリプト。
+**`networks/convert_hunyuan_image_lora_to_comfy.py`**
+
+LoRAモデルをComfyUI互換形式に変換するスクリプト。わずかに形式が異なるため、変換が必要です。以下の指定で、sd-scriptsの形式からComfyUI形式に変換できます。
+
+```bash
+python networks/convert_hunyuan_image_lora_to_comfy.py path/to/source.safetensors path/to/destination.safetensors
+```
+
+`--reverse`オプションを付けると、逆変換（ComfyUI形式からsd-scripts形式）も可能です。ただし、逆変換ができるのはこのスクリプトで変換したLoRAに限ります。他の学習ツールで作成したLoRAは変換できません。
 
 </details>
 
diff --git a/networks/convert_hunyuan_image_lora_to_comfy.py b/networks/convert_hunyuan_image_lora_to_comfy.py
index 65da2da4..df12897d 100644
--- a/networks/convert_hunyuan_image_lora_to_comfy.py
+++ b/networks/convert_hunyuan_image_lora_to_comfy.py
@@ -24,28 +24,47 @@ def main(args):
 
     logger.info(f"Converting...")
 
+    # Key mapping tables: (sd-scripts format, ComfyUI format)
+    double_blocks_mappings = [
+        ("img_mlp_fc1", "img_mlp_0"),
+        ("img_mlp_fc2", "img_mlp_2"),
+        ("img_mod_linear", "img_mod_lin"),
+        ("txt_mlp_fc1", "txt_mlp_0"),
+        ("txt_mlp_fc2", "txt_mlp_2"),
+        ("txt_mod_linear", "txt_mod_lin"),
+    ]
+
+    single_blocks_mappings = [
+        ("modulation_linear", "modulation_lin"),
+    ]
+
     keys = list(state_dict.keys())
     count = 0
+
     for k in keys:
+        new_k = k
+
         if "double_blocks" in k:
-            new_k = (
-                k.replace("img_mlp_fc1", "img_mlp_0").replace("img_mlp_fc2", "img_mlp_2").replace("img_mod_linear", "img_mod_lin")
-            )
-            new_k = (
-                new_k.replace("txt_mlp_fc1", "txt_mlp_0")
-                .replace("txt_mlp_fc2", "txt_mlp_2")
-                .replace("txt_mod_linear", "txt_mod_lin")
-            )
-            if new_k != k:
-                state_dict[new_k] = state_dict.pop(k)
-                count += 1
-                # print(f"Renamed {k} to {new_k}")
+            mappings = double_blocks_mappings
         elif "single_blocks" in k:
-            new_k = k.replace("modulation_linear", "modulation_lin")
-            if new_k != k:
-                state_dict[new_k] = state_dict.pop(k)
-                count += 1
-                # print(f"Renamed {k} to {new_k}")
+            mappings = single_blocks_mappings
+        else:
+            continue
+
+        # Apply mappings based on conversion direction
+        for src_key, dst_key in mappings:
+            if args.reverse:
+                # ComfyUI to sd-scripts: swap src and dst
+                new_k = new_k.replace(dst_key, src_key)
+            else:
+                # sd-scripts to ComfyUI: use as-is
+                new_k = new_k.replace(src_key, dst_key)
+
+        if new_k != k:
+            state_dict[new_k] = state_dict.pop(k)
+            count += 1
+            # print(f"Renamed {k} to {new_k}")
+
     logger.info(f"Converted {count} keys")
 
     # Calculate hash
@@ -64,5 +83,6 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Convert LoRA format")
     parser.add_argument("src_path", type=str, default=None, help="source path, sd-scripts format")
     parser.add_argument("dst_path", type=str, default=None, help="destination path, ComfyUI format")
+    parser.add_argument("--reverse", action="store_true", help="reverse conversion direction")
     args = parser.parse_args()
     main(args)