Refactor app.py to improve regex replacements and UI components

Enhance the regex replacement logic in several functions to use lambda functions for better readability and maintainability. Update the UI to include a configuration-only generation button, allowing users to generate metadata without starting the training process. This improves user experience by providing clearer options and streamlining the configuration workflow.

app.py

@@ -12,7 +12,7 @@ import json
 
 import gradio as gr
 import importlib
-import spaces
+# import spaces
 
 # Local modules
 from download_qwen_image_models import download_all_models, DEFAULT_MODELS_DIR
@@ -164,12 +164,14 @@ def _sync_dataset_config_jsonl(path: str, output_json: str) -> None:
     base = os.path.dirname(path)
     cache = os.path.join(base, "cache").replace("\\", "/")
 
-    new = re.sub(r"(?m)^\s*image_jsonl_file\s*=.*$", f'image_jsonl_file = "{output_json}"', txt)
+    image_line = f'image_jsonl_file = "{output_json}"'
+    new = re.sub(r"(?m)^\s*image_jsonl_file\s*=.*$", lambda _m, r=image_line: r, txt)
     if new == txt and "image_jsonl_file" not in txt:
         new = txt.rstrip("\n") + f'\nimage_jsonl_file = "{output_json}"\n'
 
     if re.search(r"(?m)^\s*cache_directory\s*=", new):
-        new = re.sub(r"(?m)^\s*cache_directory\s*=.*$", f'cache_directory = "{cache}"', new)
+        cache_line = f'cache_directory = "{cache}"'
+        new = re.sub(r"(?m)^\s*cache_directory\s*=.*$", lambda _m, r=cache_line: r, new)
     else:
         new = new.rstrip("\n") + f'\ncache_directory = "{cache}"\n'
 
@@ -414,7 +416,7 @@ def _prepare_script(
         )
 
     for pat, val in replacements.items():
-        txt = re.sub(pat, val, txt, flags=re.MULTILINE)
+        txt = re.sub(pat, lambda _m, v=val: v, txt, flags=re.MULTILINE)
 
     # Inject CONTROL_FOLDER_i if provided (uncomment/override or append)
     for i in range(8):
@@ -442,9 +444,10 @@ def _prepare_script(
 
     # Point model paths to the selected models_root
     def _replace_model_path(txt: str, key: str, rel: str) -> str:
+        repl = f"--{key} \"{models_root.rstrip('/')}/{rel}\""
         return re.sub(
             rf"--{key} \"[^\"]+\"",
-            f"--{key} \"{models_root.rstrip('/')}/{rel}\"",
+            lambda _m, r=repl: r,
             txt,
         )
 
@@ -455,9 +458,19 @@ def _prepare_script(
     txt = _replace_model_path(txt, "dit", f"dit/{dit_filename}")
 
     # Replace working dir for metadata generation to runtime /auto
-    txt = re.sub(r"^cd\s+/workspace/auto\s*$", f"cd {AUTO_DIR_RUNTIME}", txt, flags=re.MULTILINE)
+    txt = re.sub(
+        r"^cd\s+/workspace/auto\s*$",
+        lambda _m: f"cd {AUTO_DIR_RUNTIME}",
+        txt,
+        flags=re.MULTILINE,
+    )
     # Ensure musubi-tuner path matches runtime location
-    txt = re.sub(r"^cd\s+/musubi-tuner\s*$", f"cd {re.escape(MUSUBI_TUNER_DIR_RUNTIME)}", txt, flags=re.MULTILINE)
+    txt = re.sub(
+        r"^cd\s+/musubi-tuner\s*$",
+        lambda _m: f"cd {MUSUBI_TUNER_DIR_RUNTIME}",
+        txt,
+        flags=re.MULTILINE,
+    )
 
     # ZeroGPU compatibility: avoid spawning via 'accelerate launch'.
     # Run the training module directly in-process so GPU stays attached
@@ -477,7 +490,8 @@ def _prepare_script(
         txt = re.sub(r"--save_every_n_epochs\s+\d+",
                      f"--save_every_n_epochs {override_save_every}", txt)
     if override_run_name:
-        txt = re.sub(r"^RUN_NAME=.*$", f"RUN_NAME={_bash_quote(override_run_name)}", txt, flags=re.MULTILINE)
+        repl = f"RUN_NAME={_bash_quote(override_run_name)}"
+        txt = re.sub(r"^RUN_NAME=.*$", lambda _m, r=repl: r, txt, flags=re.MULTILINE)
 
     # Inject prefix/suffix flags for metadata creation
     extra_lines: List[str] = []
@@ -526,7 +540,7 @@ def _prepare_script(
         pattern = rf"(?m)^\s*{name}\s*=.*$"
         replacement = f'{name}="{value}"' if not str(value).isdigit() else f'{name}={value}'
         if re.search(pattern, txt):
-            txt = re.sub(pattern, replacement, txt)
+            txt = re.sub(pattern, lambda _m, r=replacement: r, txt)
         else:
             txt = f"{replacement}\n" + txt
 
@@ -655,7 +669,7 @@ def _startup_install_musubi_deps() -> None:
             print("[QIE] WARN: musubi-tuner installation failed. Continuing.")
 
 
-@spaces.GPU
+# @spaces.GPU
 def run_training(
     output_name: str,
     caption: str,
@@ -723,8 +737,8 @@ def run_training(
     # Auto-generate dataset directory name
     import time
     ds_name = f"dataset_{int(time.time())}"
-    ds_dir = os.path.join(DATA_ROOT_RUNTIME, ds_name)
-    run_out_dir = os.path.join(ds_dir, output_name.strip())
+    ds_dir = os.path.abspath(os.path.join(DATA_ROOT_RUNTIME, ds_name))
+    run_out_dir = os.path.abspath(os.path.join(ds_dir, output_name.strip()))
     img_folder_name = DEFAULT_IMAGE_FOLDER
     img_dir = os.path.join(ds_dir, img_folder_name)
     os.makedirs(img_dir, exist_ok=True)
@@ -1013,8 +1027,6 @@ def build_ui() -> gr.Blocks:
                             cr_w = gr.Number(label="Control resolution W", value=1024, precision=0)
                             cr_h = gr.Number(label="Control resolution H", value=1024, precision=0)
                             te_bs = gr.Number(label="TE cache batch size", value=16, precision=0)
-                        with gr.Row():
-                            config_only = gr.Checkbox(label="設定のみ生成 (学習しない)", value=False)
 
                 with gr.Accordion("Target Image", elem_classes=["pad-section_0"]):
                     with gr.Group():
@@ -1112,11 +1124,13 @@ def build_ui() -> gr.Blocks:
 
                 # Models root / OUTPUT_DIR_BASE / DATASET_CONFIG are auto-resolved at runtime; no user input needed.
 
-                run_btn = gr.Button("Start Training", variant="primary")
-                run_out_dir_box = gr.Textbox(label="今回の出力フォルダ", lines=1, interactive=False)
+                with gr.Row():
+                    run_btn = gr.Button("Start Training", variant="primary")
+                    config_btn = gr.Button("設定のみ生成", variant="secondary")
+                run_out_dir_box = gr.Textbox(label="出力フォルダ", lines=1, interactive=False)
                 scripts_files = gr.Files(label="Scripts & Config (live)", interactive=False)
-                logs = gr.Textbox(label="Logs", lines=20)
                 ckpt_files = gr.Files(label="Checkpoints (live)", interactive=False)
+                logs = gr.Textbox(label="Logs", lines=20)
                 with gr.Row():
                     refresh_scripts_btn = gr.Button("ファイルを再取得", variant="secondary")
 
@@ -1133,6 +1147,9 @@ def build_ui() -> gr.Blocks:
                 ctrl6_files.change(fn=_files_to_gallery, inputs=ctrl6_files, outputs=ctrl6_gallery)
                 ctrl7_files.change(fn=_files_to_gallery, inputs=ctrl7_files, outputs=ctrl7_gallery)
 
+                config_only_off = gr.State(False)
+                config_only_on = gr.State(True)
+
                 run_btn.click(
                     fn=run_training,
                     inputs=[
@@ -1147,7 +1164,25 @@ def build_ui() -> gr.Blocks:
                         ctrl7_files, ctrl7_prefix, ctrl7_suffix,
                         lr_input, dim_input,
                         tr_w, tr_h, train_bs, cr_w, cr_h, te_bs,
-                        seed_input, max_epochs, save_every, config_only,
+                        seed_input, max_epochs, save_every, config_only_off,
+                    ],
+                    outputs=[logs, ckpt_files, scripts_files, run_out_dir_box],
+                )
+                config_btn.click(
+                    fn=run_training,
+                    inputs=[
+                        output_name, caption, images_input, main_prefix, main_suffix,
+                        ctrl0_files, ctrl0_prefix, ctrl0_suffix,
+                        ctrl1_files, ctrl1_prefix, ctrl1_suffix,
+                        ctrl2_files, ctrl2_prefix, ctrl2_suffix,
+                        ctrl3_files, ctrl3_prefix, ctrl3_suffix,
+                        ctrl4_files, ctrl4_prefix, ctrl4_suffix,
+                        ctrl5_files, ctrl5_prefix, ctrl5_suffix,
+                        ctrl6_files, ctrl6_prefix, ctrl6_suffix,
+                        ctrl7_files, ctrl7_prefix, ctrl7_suffix,
+                        lr_input, dim_input,
+                        tr_w, tr_h, train_bs, cr_w, cr_h, te_bs,
+                        seed_input, max_epochs, save_every, config_only_on,
                     ],
                     outputs=[logs, ckpt_files, scripts_files, run_out_dir_box],
                 )