Refactor app.py to update prefix/suffix naming conventions for metadata creation and enhance UI with new training hyperparameter inputs. Modify train_QIE.sh to utilize dynamic hyperparameter values for training execution, improving configurability and user experience.

app.py

@@ -322,9 +322,9 @@ def _prepare_script(
     # Inject prefix/suffix flags for metadata creation
     extra_lines: List[str] = []
     if (target_prefix or ""):
-        extra_lines.append(f"  --target_prefix {_bash_quote(target_prefix)} \\")
+        extra_lines.append(f"  --main_prefix {_bash_quote(target_prefix)} \\")
     if (target_suffix or ""):
-        extra_lines.append(f"  --target_suffix {_bash_quote(target_suffix)} \\")
+        extra_lines.append(f"  --main_suffix {_bash_quote(target_suffix)} \\")
     for i in range(8):
         pre = control_prefixes[i] if (control_prefixes and i < len(control_prefixes)) else None
         suf = control_suffixes[i] if (control_suffixes and i < len(control_suffixes)) else None
@@ -489,6 +489,9 @@ def run_training(
     control7_uploads: Any,
     ctrl7_prefix: str,
     ctrl7_suffix: str,
+    learning_rate: str,
+    network_dim: int,
+    seed: int,
     max_epochs: int,
     save_every: int,
 ) -> Iterable[tuple]:
@@ -593,6 +596,9 @@ def run_training(
         target_suffix=(target_suffix or ""),
         control_prefixes=[ctrl0_prefix, ctrl1_prefix, ctrl2_prefix, ctrl3_prefix, ctrl4_prefix, ctrl5_prefix, ctrl6_prefix, ctrl7_prefix],
         control_suffixes=[ctrl0_suffix, ctrl1_suffix, ctrl2_suffix, ctrl3_suffix, ctrl4_suffix, ctrl5_suffix, ctrl6_suffix, ctrl7_suffix],
+        override_learning_rate=(learning_rate or None),
+        override_network_dim=int(network_dim) if network_dim is not None else None,
+        override_seed=int(seed) if seed is not None else None,
     )
 
 
@@ -653,6 +659,14 @@ def build_ui() -> gr.Blocks:
             output_name = gr.Textbox(label="OUTPUT NAME", placeholder="my_lora_output", lines=1)
             caption = gr.Textbox(label="CAPTION", placeholder="A photo of ...", lines=2)
 
+        # Training options near OUTPUT NAME
+        with gr.Row():
+            lr_input = gr.Textbox(label="Learning rate", value="1e-3")
+            dim_input = gr.Number(label="Network dim", value=4, precision=0)
+            seed_input = gr.Number(label="Seed", value=42, precision=0)
+            max_epochs = gr.Number(label="Max epochs", value=100, precision=0)
+            save_every = gr.Number(label="Save every N epochs", value=10, precision=0)
+
         with gr.Row():
             with gr.Column(scale=3):
                 images_input = gr.File(label="Upload target images", file_count="multiple", type="filepath")
@@ -728,12 +742,10 @@ def build_ui() -> gr.Blocks:
         logs = gr.Textbox(label="Logs", lines=20)
         ckpt_files = gr.Files(label="Checkpoints (live)", interactive=False)
 
-        with gr.Row():
-            max_epochs = gr.Number(label="Max epochs (this run)", value=10, precision=0)
-            save_every = gr.Number(label="Save every N epochs", value=5, precision=0)
+        # moved max_epochs/save_every above next to OUTPUT NAME
 
         # Wire previews
-        images_input.change(fn=_files_to_gallery, inputs=images_input, outputs=target_gallery)
+        images_input.change(fn=_files_to_gallery, inputs=images_input, outputs=main_gallery)
         ctrl0_files.change(fn=_files_to_gallery, inputs=ctrl0_files, outputs=ctrl0_gallery)
         ctrl1_files.change(fn=_files_to_gallery, inputs=ctrl1_files, outputs=ctrl1_gallery)
         ctrl2_files.change(fn=_files_to_gallery, inputs=ctrl2_files, outputs=ctrl2_gallery)
@@ -746,7 +758,7 @@ def build_ui() -> gr.Blocks:
         run_btn.click(
             fn=run_training,
             inputs=[
-                output_name, caption, images_input, target_prefix, target_suffix,
+                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,
@@ -755,7 +767,7 @@ def build_ui() -> gr.Blocks:
                 ctrl5_files, ctrl5_prefix, ctrl5_suffix,
                 ctrl6_files, ctrl6_prefix, ctrl6_suffix,
                 ctrl7_files, ctrl7_prefix, ctrl7_suffix,
-                max_epochs, save_every,
+                lr_input, dim_input, seed_input, max_epochs, save_every,
             ],
             outputs=[logs, ckpt_files],
         )

train_QIE.sh

@@ -39,6 +39,13 @@ OUTPUT_DIR_BASE="/workspace/auto/train_LoRA"
 DATASET_CONFIG="/workspace/auto/dataset_QIE.toml"
 OUTPUT_JSON="${DATASET_DIR%/}/metadata.jsonl"
 
+# Training hyperparameters (can be overridden by app)
+LEARNING_RATE="1e-3"
+NETWORK_DIM=4
+SEED=42
+MAX_TRAIN_EPOCHS=100
+SAVE_EVERY_N_EPOCHS=10
+
 # Build control args from folder names with auto-detect fallback
 CONTROL_ARGS=()
 for i in {0..7}; do
@@ -121,15 +128,15 @@ accelerate launch src/musubi_tuner/qwen_image_train_network.py \
 --weighting_scheme none \
 --discrete_flow_shift 2.0  \
 --optimizer_type adamw8bit \
---learning_rate 1e-3 \
+--learning_rate "$LEARNING_RATE" \
 --gradient_checkpointing \
 --max_data_loader_n_workers 2 \
 --persistent_data_loader_workers \
 --network_module networks.lora_qwen_image \
---network_dim 4 \
---max_train_epochs 100 \
---save_every_n_epochs 10 \
---seed 42 \
+--network_dim "$NETWORK_DIM" \
+--max_train_epochs "$MAX_TRAIN_EPOCHS" \
+--save_every_n_epochs "$SAVE_EVERY_N_EPOCHS" \
+--seed "$SEED" \
 --output_dir "${OUTPUT_DIR_BASE}/${RUN_NAME}" \
 --output_name "${RUN_NAME}" \
 --ddp_gradient_as_bucket_view \