Update app.py

app.py

@@ -1,16 +1,6 @@
 # =====================================================================
 # ForgeCaptions - Gradio app for single & batch image captioning
 # =====================================================================
-# CHANGELOG (this version)
-# - GPU-safe: all CUDA only inside @spaces.GPU functions.
-# - Restored: Single tab + Batch chunking (Auto / All-at-once / Manual step).
-# - Shape Aliases: supports comma/pipe-separated synonyms per row.
-# - Default caption style: "Character training (long)".
-# - Model Instructions + Caption Style in minimizable accordions.
-# - Excel export: thumbnail size slider controls image scaling & row height.
-# - Header logo scaled to the full text stack (centered).
-# - Kept gallery & table positions unchanged; scroll sync retained.
-# =====================================================================
 
 # ------------------------------
 # 0) Imports & environment
@@ -408,22 +398,30 @@ def run_batch(
     top_p: float,
     max_tokens: int,
     max_side: int,
-) -> Tuple[List[dict], list, list, str]:
-    """
-    Process a list of file paths and append results to session_rows.
-    Returns: updated rows, gallery_pairs, table_rows, status_text
+    time_budget_s: float | None = None,                    # respects Zero-GPU window
+    progress: gr.Progress = gr.Progress(track_tqdm=True),  # drives the progress bar
+) -> Tuple[List[dict], list, list, str, List[str], int, int]:
+     """
+    Returns:
+      session_rows, gallery_pairs, table_rows, status_text,
+      leftover_files, processed_in_this_call, total_in_this_call
     """
     session_rows = session_rows or []
-    files = files or []
-    if not files:
+    files = [f for f in (files or []) if f and os.path.exists(f)]
+    total = len(files)
+    processed = 0
+
+    if total == 0:
         gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
                          for r in session_rows if (r.get("thumb_path") or r.get("path"))]
         table_rows = [[r.get("filename",""), r.get("caption","")] for r in session_rows]
-        return session_rows, gallery_pairs, table_rows, f"Saved • {time.strftime('%H:%M:%S')}"
+        return session_rows, gallery_pairs, table_rows, f"Saved • {time.strftime('%H:%M:%S')}", [], 0, 0
 
-    for path in files:
-        if not path or not os.path.exists(path):
-            continue
+    start = time.time()
+    leftover: List[str] = []
+    
+    # Progress bar shows inside the GPU worker
+    for idx, path in enumerate(progress.tqdm(files, desc="Captioning")):
         try:
             im = Image.open(path).convert("RGB")
         except Exception:
@@ -436,12 +434,26 @@ def run_batch(
         filename = os.path.basename(path)
         thumb = ensure_thumb(path, 256)
         session_rows.append({"filename": filename, "caption": cap, "path": path, "thumb_path": thumb})
+        processed += 1
+
+       # Time-slice to avoid Zero GPU timeouts
+        if time_budget_s and (time.time() - start) >= float(time_budget_s):
+            leftover = files[idx+1:]
+            break 
 
     save_session(session_rows)
     gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
                      for r in session_rows if (r.get("thumb_path") or r.get("path"))]
     table_rows = [[r.get("filename",""), r.get("caption","")] for r in session_rows]
-    return session_rows, gallery_pairs, table_rows, f"Saved • {time.strftime('%H:%M:%S')}"
+    return (
+        session_rows,
+        gallery_pairs,
+        table_rows,
+        f"Saved • {time.strftime('%H:%M:%S')}",
+        leftover,
+        processed,
+        total,
+    )
 
 # Ensure Spaces detects at least one GPU function at startup
 @gpu
@@ -632,6 +644,7 @@ setTimeout(() => {{
                     label="Batch mode"
                 )
                 chunk_size = gr.Slider(1, 50, value=10, step=1, label="Chunk size")
+                gpu_budget = gr.Slider(20, 110, value=55, step=5, label="Max seconds per GPU call")
 
     # -- Keep instruction text in sync with controls and persist to settings
     def _refresh_instruction(styles, extra, name_value, trigv, begv, endv, excel_px, ms):
@@ -676,6 +689,7 @@ setTimeout(() => {{
     # ---- Results (UNCHANGED POSITION): Gallery left, Table right
     rows_state  = gr.State(load_session())
     autosave_md = gr.Markdown("Ready.")
+    progress_md  = gr.Markdown("") 
     remaining_state = gr.State([])  # for manual step mode
 
     with gr.Row():
@@ -770,45 +784,66 @@ setTimeout(() => {{
         s = load_settings()
         return s.get("temperature", 0.6), s.get("top_p", 0.9), s.get("max_tokens", 256)
 
-    def _run_click(files, rows, instr, ms, mode, csize):
+    def _run_click(files, rows, instr, ms, mode, csize, budget_s):
         t, p, m = _tpms()
         files = files or []
+        
         # Manual step → process first chunk only
         if mode == "Manual (step)" and files:
             chunks = _split_chunks(files, int(csize))
             batch = chunks[0]
             remaining = sum(chunks[1:], [])
-            new_rows, gal, tbl, stamp = run_batch(batch, rows or [], instr, t, p, m, int(ms))
+            new_rows, gal, tbl, stamp, leftover_from_batch, done, total = run_batch(
+                batch, rows or [], instr, t, p, m, int(ms), float(budget_s)
+            )
+            remaining = (leftover_from_batch or []) + remaining
             panel_vis = gr.update(visible=bool(remaining))
             msg = f"{len(remaining)} files remain. Process next chunk?"
-            return new_rows, gal, tbl, stamp, remaining, panel_vis, gr.update(value=msg)
-        # Auto / all-at-once → process everything in one go
-        else:
-            new_rows, gal, tbl, stamp = run_batch(files, rows or [], instr, t, p, m, int(ms))
-            return new_rows, gal, tbl, stamp, [], gr.update(visible=False), gr.update(value="")
+            prog = f"Batch progress: {done}/{total} processed in this step • Remaining overall: {len(remaining)}"
+            return new_rows, gal, tbl, stamp, remaining, panel_vis, gr.update(value=msg), gr.update(value=prog)
 
+         # Auto / All-at-once (still obey time budget)
+        new_rows, gal, tbl, stamp, leftover, done, total = run_batch(
+            files, rows or [], instr, t, p, m, int(ms), float(budget_s)
+        )
+        panel_vis = gr.update(visible=bool(leftover))
+        msg = f"{len(leftover)} files remain. Process next chunk?" if leftover else ""
+        prog = f"Batch progress: {done}/{total} processed in this call • Remaining: {len(leftover)}"
+        return new_rows, gal, tbl, stamp, leftover, panel_vis, gr.update(value=msg), gr.update(value=prog)
+
+        
     run_button.click(
         _run_click,
-        inputs=[input_files, rows_state, instruction_preview, max_side, chunk_mode, chunk_size],
-        outputs=[rows_state, gallery, table, autosave_md, remaining_state, step_panel, step_msg]
+        inputs=[input_files, rows_state, instruction_preview, max_side, chunk_mode, chunk_size, gpu_budget],
+        outputs=[rows_state, gallery, table, autosave_md, remaining_state, step_panel, step_msg, progress_md]
     )
 
-    def _step_next(remain, rows, instr, ms, csize):
+    def _step_next(remain, rows, instr, ms, csize, budget_s):
         t, p, m = _tpms()
         remain = remain or []
         if not remain:
-            return rows, gr.update(value="No files remaining."), gr.update(visible=False), [], [], [], "Saved."
+            return rows, gr.update(value="No files remaining."), gr.update(visible=False), [], [], [], "Saved.", gr.update(value="")
+    
         batch = remain[:int(csize)]
         leftover = remain[int(csize):]
-        new_rows, gal, tbl, stamp = run_batch(batch, rows or [], instr, t, p, m, int(ms))
+        new_rows, gal, tbl, stamp, leftover_from_batch, done, total = run_batch(
+            batch, rows or [], instr, t, p, m, int(ms), float(budget_s)
+        )
+        leftover = (leftover_from_batch or []) + leftover
         panel_vis = gr.update(visible=bool(leftover))
         msg = f"{len(leftover)} files remain. Process next chunk?" if leftover else "All done."
-        return new_rows, msg, panel_vis, leftover, gal, tbl, stamp
-
+        prog = f"Batch progress: {done}/{total} processed in this step • Remaining overall: {len(leftover)}"
+        return new_rows, msg, panel_vis, leftover, gal, tbl, stamp, gr.update(value=prog)
+    
+    # hookup (with gpu_budget slider)
     step_next.click(
         _step_next,
-        inputs=[remaining_state, rows_state, instruction_preview, max_side, chunk_size],
-        outputs=[rows_state, step_msg, step_panel, remaining_state, gallery, table, autosave_md]
+        inputs=[remaining_state, rows_state, instruction_preview, max_side, chunk_size, gpu_budget],
+        outputs=[rows_state, step_msg, step_panel, remaining_state, gallery, table, autosave_md, progress_md]
+)
+
+# if you don't have gpu_budget: wrap with a lambda and pass 55 as above
+
     )
 
     def _step_finish():