Update app.py

app.py

@@ -359,9 +359,11 @@ def _build_inputs(im: Image.Image, instr: str, dtype) -> Dict[str, Any]:
         inputs["pixel_values"] = inputs["pixel_values"].to(dtype)
     return inputs
 
-@gpu
 @torch.no_grad()
 def caption_once(im: Image.Image, instr: str, temp: float, top_p: float, max_tokens: int) -> str:
+    """
+    NOTE: Not GPU-decorated on purpose; call this only from within a @gpu function.
+    """
     model, device, dtype = get_model()
     inputs = _build_inputs(im, instr, dtype)
     inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
@@ -401,7 +403,7 @@ def run_batch(
     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
@@ -419,7 +421,7 @@ def run_batch(
 
     start = time.time()
     leftover: List[str] = []
-    
+
     # Progress bar shows inside the GPU worker
     for idx, path in enumerate(progress.tqdm(files, desc="Captioning")):
         try:
@@ -436,10 +438,10 @@ def run_batch(
         session_rows.append({"filename": filename, "caption": cap, "path": path, "thumb_path": thumb})
         processed += 1
 
-       # Time-slice to avoid Zero GPU timeouts
+        # Time-slice to avoid Zero GPU timeouts
         if time_budget_s and (time.time() - start) >= float(time_budget_s):
             leftover = files[idx+1:]
-            break 
+            break
 
     save_session(session_rows)
     gallery_pairs = [((r.get("thumb_path") or r.get("path")), r.get("caption",""))
@@ -455,7 +457,7 @@ def run_batch(
         total,
     )
 
-# Ensure Spaces detects at least one GPU function at startup
+# Ensure Spaces detects at least one GPU function at startup (without CUDA in main proc)
 @gpu
 @torch.no_grad()
 def _gpu_startup_warm():
@@ -689,7 +691,7 @@ setTimeout(() => {{
     # ---- Results (UNCHANGED POSITION): Gallery left, Table right
     rows_state  = gr.State(load_session())
     autosave_md = gr.Markdown("Ready.")
-    progress_md  = gr.Markdown("") 
+    progress_md = gr.Markdown("") 
     remaining_state = gr.State([])  # for manual step mode
 
     with gr.Row():
@@ -802,7 +804,7 @@ setTimeout(() => {{
             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)
+        # 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)
         )
@@ -811,7 +813,6 @@ setTimeout(() => {{
         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, gpu_budget],
@@ -835,15 +836,10 @@ setTimeout(() => {{
         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, 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():
@@ -884,7 +880,6 @@ setTimeout(() => {{
         alias_table = gr.Dataframe(
             headers=["shape (token or synonyms)", "name to insert"],
             value=init_rows,
-            col_count=(2, "fixed"),
             row_count=(max(1, len(init_rows)), "dynamic"),
             datatype=["str","str"],
             type="array",