Refactor run_training function in app.py to return tuples for error and log messages, enhancing clarity in output handling. Update UI to remove user input for models_root, output_dir_base, and dataset_config, which are now resolved at runtime. Implement logic to locate the latest LoRA file for download after training completion.

app.py

@@ -346,18 +346,15 @@ def run_training(
     control5_uploads: Any,
     control6_uploads: Any,
     control7_uploads: Any,
-    models_root: str,
-    output_dir_base: str,
-    dataset_config: str,
     max_epochs: int,
     save_every: int,
-) -> Iterable[str]:
+) -> Iterable[tuple]:
     # Basic validation
     if not output_name.strip():
-        yield "[ERROR] OUTPUT NAME is required."
+        yield ("[ERROR] OUTPUT NAME is required.", None)
         return
     if not caption.strip():
-        yield "[ERROR] CAPTION is required."
+        yield ("[ERROR] CAPTION is required.", None)
         return
 
     # Ensure /auto holds helper files expected by the script
@@ -376,10 +373,10 @@ def run_training(
     # Ingest uploads into dataset folders
     base_files = _extract_paths(image_uploads)
     if not base_files:
-        yield "[ERROR] No images uploaded for IMAGE_FOLDER."
+        yield ("[ERROR] No images uploaded for IMAGE_FOLDER.", None)
         return
     base_filenames = _copy_uploads(base_files, img_dir)
-    yield f"[QIE] Copied {len(base_filenames)} base images to {img_dir}"
+    yield (f"[QIE] Copied {len(base_filenames)} base images to {img_dir}", None)
 
     # Prepare control sets
     control_upload_sets = [
@@ -408,14 +405,14 @@ def run_training(
             replicated = uploads * len(base_filenames)
             _copy_uploads(replicated, cdir, rename_to=base_filenames)
         else:
-            yield f"[ERROR] control_{i}: file count {len(uploads)} must be 1 or {len(base_filenames)}."
+            yield (f"[ERROR] control_{i}: file count {len(uploads)} must be 1 or {len(base_filenames)}.", None)
             return
         control_dirs.append(folder_name)
         any_control = True
-        yield f"[QIE] Copied {len(uploads)} control_{i} images to {cdir}"
+        yield (f"[QIE] Copied {len(uploads)} control_{i} images to {cdir}", None)
 
     if not any_control:
-        yield "[ERROR] At least one control folder is required for edit-plus training."
+        yield ("[ERROR] At least one control folder is required for edit-plus training.", None)
         return
 
     # Prepare script with user parameters
@@ -425,9 +422,15 @@ def run_training(
     ]
 
     # Decide dataset_config path with fallback to runtime auto dir
-    ds_conf_input = (dataset_config or "").strip()
-    ds_conf_runtime = str(Path(AUTO_DIR_RUNTIME) / "dataset_QIE.toml")
-    ds_conf = ds_conf_input if (ds_conf_input and os.path.exists(ds_conf_input)) else ds_conf_runtime
+    ds_conf = str(Path(AUTO_DIR_RUNTIME) / "dataset_QIE.toml")
+
+    # Resolve models_root and output_dir_base at runtime
+    models_root = MODELS_ROOT_RUNTIME
+    out_base = os.path.join(AUTO_DIR_RUNTIME, "train_LoRA")
+    try:
+        os.makedirs(out_base, exist_ok=True)
+    except Exception:
+        pass
 
     tmp_script = _prepare_script(
         dataset_name=ds_name,
@@ -435,8 +438,8 @@ def run_training(
         data_root=DATA_ROOT_RUNTIME,
         image_folder=img_folder_name,
         control_folders=control_folders,
-        models_root=models_root.strip() or MODELS_ROOT_RUNTIME,
-        output_dir_base=(output_dir_base.strip() or None),
+        models_root=models_root,
+        output_dir_base=out_base,
         dataset_config=ds_conf,
         override_max_epochs=max_epochs if max_epochs and max_epochs > 0 else None,
         override_save_every=save_every if save_every and save_every > 0 else None,
@@ -445,8 +448,8 @@ def run_training(
 
 
     shell = _pick_shell()
-    yield f"[QIE] Using shell: {shell}"
-    yield f"[QIE] Running script: {tmp_script}"
+    yield (f"[QIE] Using shell: {shell}", None)
+    yield (f"[QIE] Running script: {tmp_script}", None)
 
     # Run and stream output
     proc = subprocess.Popen(
@@ -460,10 +463,26 @@ def run_training(
     try:
         assert proc.stdout is not None
         for line in proc.stdout:
-            yield line.rstrip("\n")
+            yield (line.rstrip("\n"), None)
     finally:
         code = proc.wait()
-        yield f"[QIE] Exit code: {code}"
+        # Try to locate latest LoRA file for download
+        lora_path = None
+        try:
+            out_dir = os.path.join(out_base, output_name.strip())
+            if os.path.isdir(out_dir):
+                cand = []
+                for root, _, files in os.walk(out_dir):
+                    for fn in files:
+                        if fn.lower().endswith(".safetensors"):
+                            full = os.path.join(root, fn)
+                            cand.append((os.path.getmtime(full), full))
+                if cand:
+                    cand.sort()
+                    lora_path = cand[-1][1]
+        except Exception:
+            pass
+        yield (f"[QIE] Exit code: {code}", lora_path)
 
 
 def build_ui() -> gr.Blocks:
@@ -500,13 +519,11 @@ def build_ui() -> gr.Blocks:
             with gr.Row():
                 ctrl7_files = gr.File(label="Upload control_7 images", file_count="multiple", type="filepath")
 
-        with gr.Row():
-            models_root = gr.Textbox(label="Models root", value=MODELS_ROOT_RUNTIME)
-            output_dir_base = gr.Textbox(label="OUTPUT_DIR_BASE", value=DEFAULT_OUTPUT_DIR_BASE)
-            dataset_config = gr.Textbox(label="DATASET_CONFIG", value=str(Path(AUTO_DIR_RUNTIME) / "dataset_QIE.toml"))
+        # Models root / OUTPUT_DIR_BASE / DATASET_CONFIG are auto-resolved at runtime; no user input needed.
 
         run_btn = gr.Button("Start Training", variant="primary")
         logs = gr.Textbox(label="Logs", lines=20)
+        lora_file = gr.File(label="Download LoRA", interactive=False)
 
         with gr.Row():
             max_epochs = gr.Number(label="Max epochs (this run)", value=10, precision=0)
@@ -517,10 +534,9 @@ def build_ui() -> gr.Blocks:
             inputs=[
                 output_name, caption, images_input,
                 ctrl0_files, ctrl1_files, ctrl2_files, ctrl3_files, ctrl4_files, ctrl5_files, ctrl6_files, ctrl7_files,
-                models_root, output_dir_base, dataset_config,
                 max_epochs, save_every,
             ],
-            outputs=logs,
+            outputs=[logs, lora_file],
         )
 
     return demo