Add checkpoint listing functionality in app.py to track and display model checkpoints during training. Update run_training to yield checkpoint information and enhance Gradio UI with checkpoint file outputs for improved user experience.

app.py

@@ -7,6 +7,7 @@ import sys
 import tempfile
 from pathlib import Path
 from typing import Dict, Iterable, List, Optional, Any, Tuple
+import time
 import json
 
 import gradio as gr
@@ -173,6 +174,25 @@ def _copy_uploads(
     return used_names
 
 
+def _list_checkpoints(out_dir: str, limit: int = 20) -> List[str]:
+    try:
+        if not out_dir or not os.path.isdir(out_dir):
+            return []
+        items: List[Tuple[float, str]] = []
+        for root, _, files in os.walk(out_dir):
+            for fn in files:
+                if fn.lower().endswith('.safetensors'):
+                    full = os.path.join(root, fn)
+                    try:
+                        items.append((os.path.getmtime(full), full))
+                    except Exception:
+                        pass
+        items.sort(reverse=True)
+        return [p for _, p in items[:limit]]
+    except Exception:
+        return []
+
+
 def _prepare_script(
     dataset_name: str,
     caption: str,
@@ -445,11 +465,14 @@ def run_training(
 ) -> Iterable[tuple]:
     # Basic validation
     log_buf = ""
+    ckpts: List[str] = []
     if not output_name.strip():
-        yield ("[ERROR] OUTPUT NAME is required.", None)
+        log_buf += "[ERROR] OUTPUT NAME is required.\n"
+        yield (log_buf, ckpts, None)
         return
     if not caption.strip():
-        yield ("[ERROR] CAPTION is required.", None)
+        log_buf += "[ERROR] CAPTION is required.\n"
+        yield (log_buf, ckpts, None)
         return
 
     # Ensure /auto holds helper files expected by the script
@@ -468,11 +491,12 @@ 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.", None)
+        log_buf += "[ERROR] No images uploaded for IMAGE_FOLDER.\n"
+        yield (log_buf, ckpts, None)
         return
     base_filenames = _copy_uploads(base_files, img_dir)
     log_buf += f"[QIE] Copied {len(base_filenames)} base images to {img_dir}\n"
-    yield (log_buf, None)
+    yield (log_buf, ckpts, None)
 
     # Prepare control sets
     control_upload_sets = [
@@ -487,7 +511,8 @@ def run_training(
     ]
     # Require control_0; others optional
     if not control_upload_sets[0]:
-        yield ("[ERROR] control_0 images are required.", None)
+        log_buf += "[ERROR] control_0 images are required.\n"
+        yield (log_buf, ckpts, None)
         return
 
     control_dirs: List[Optional[str]] = []
@@ -502,7 +527,7 @@ def run_training(
         _copy_uploads(uploads, cdir)
         control_dirs.append(folder_name)
         log_buf += f"[QIE] Copied {len(uploads)} control_{i} images to {cdir}\n"
-        yield (log_buf, None)
+        yield (log_buf, ckpts, None)
 
     # Metadata.jsonl will be generated by create_image_caption_json.py in train_QIE.sh
 
@@ -545,7 +570,9 @@ def run_training(
     shell = _pick_shell()
     log_buf += f"[QIE] Using shell: {shell}\n"
     log_buf += f"[QIE] Running script: {tmp_script}\n"
-    yield (log_buf, None)
+    out_dir = os.path.join(out_base, output_name.strip())
+    ckpts = _list_checkpoints(out_dir)
+    yield (log_buf, ckpts, None)
 
     # Run and stream output
     proc = subprocess.Popen(
@@ -558,29 +585,24 @@ def run_training(
     )
     try:
         assert proc.stdout is not None
+        i = 0
         for line in proc.stdout:
             log_buf += line
-            yield (log_buf, None)
+            i += 1
+            if i % 30 == 0:
+                ckpts = _list_checkpoints(out_dir)
+            yield (log_buf, ckpts, None)
     finally:
         code = proc.wait()
         # 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]
+            ckpts = _list_checkpoints(out_dir)
         except Exception:
             pass
+        lora_path = ckpts[0] if ckpts else None
         log_buf += f"[QIE] Exit code: {code}\n"
-        yield (log_buf, lora_path)
+        yield (log_buf, ckpts, lora_path)
 
 
 def build_ui() -> gr.Blocks:
@@ -685,7 +707,8 @@ def build_ui() -> gr.Blocks:
 
         run_btn = gr.Button("Start Training", variant="primary")
         logs = gr.Textbox(label="Logs", lines=20)
-        lora_file = gr.File(label="Download LoRA", interactive=False)
+        ckpt_files = gr.Files(label="Checkpoints (live)", interactive=False)
+        lora_file = gr.File(label="Download LoRA (latest)", interactive=False)
 
         with gr.Row():
             max_epochs = gr.Number(label="Max epochs (this run)", value=10, precision=0)
@@ -705,7 +728,7 @@ def build_ui() -> gr.Blocks:
                 ctrl7_files, ctrl7_prefix, ctrl7_suffix,
                 max_epochs, save_every,
             ],
-            outputs=[logs, lora_file],
+            outputs=[logs, ckpt_files, lora_file],
         )
 
     return demo