Try full pipeline GPU residency before falling back to cpu offload

Each ZeroGPU call runs in a fresh worker (offload hooks are always unset
at entry), so enable_model_cpu_offload bought no cross-call reuse — it
just replaced one bulk pipe.to(device) transfer with several slower,
hook-managed component transfers, adding ~55s of pure data movement to a
~62s call per the FireRed-Dev logs. Try the fast full-GPU path first and
only fall back to offload on an actual CUDA OOM (small MIG slices).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

app.py

@@ -402,13 +402,22 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
         print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
         torch.cuda.reset_peak_memory_stats()
 
-    # Sequential offload: only the component currently doing work (text_encoder,
-    # then transformer, then vae — see model_cpu_offload_seq) sits on GPU at a
-    # time, instead of the full ~37GB pipeline resident simultaneously. Needed
-    # to fit on smaller MIG slices (e.g. the 47GB 2g.48gb partition) without OOM.
+    # Each ZeroGPU call runs in a fresh worker (hooks are always unset here),
+    # so cpu_offload buys no cross-call reuse — it only trades one bulk
+    # to(device) transfer for several slower hook-managed ones. Try moving
+    # the full ~37GB pipeline to GPU in one shot first (fast); only fall back
+    # to enable_model_cpu_offload (keeps one component resident at a time) if
+    # that OOMs on a smaller MIG slice (e.g. the 47GB 2g.48gb partition).
     if getattr(pipe.transformer, "_hf_hook", None) is None:
-        pipe.enable_model_cpu_offload(device=device)
-        print(f"[infer] enabled sequential cpu offload on {device}")
+        try:
+            pipe.to(device)
+            print(f"[infer] moved full pipe to {device} — t={time.perf_counter()-t0:.1f}s")
+        except torch.cuda.OutOfMemoryError:
+            print(f"[infer] OOM moving full pipe to {device}, falling back to cpu offload")
+            pipe.to("cpu")
+            torch.cuda.empty_cache()
+            pipe.enable_model_cpu_offload(device=device)
+            print(f"[infer] enabled cpu offload on {device} (fallback)")
     print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
 
     print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
@@ -452,9 +461,11 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
         timer.print_timings()
         raise
     finally:
-        # No manual pipe.to("cpu") — that fights the offload hooks' own device
-        # bookkeeping. Hooks already return each component to CPU after its
-        # forward; empty_cache() just reclaims the freed GPU blocks.
+        # No manual pipe.to("cpu"): in the offload-fallback case that fights the
+        # hooks' own device bookkeeping (they return each component to CPU after
+        # its forward), and in the normal fast-path case the worker's GPU access
+        # is reclaimed by ZeroGPU when this call returns regardless — paying for
+        # a D2H transfer here would just be wasted GPU-billed time.
         gc.collect()
         torch.cuda.empty_cache()
         print(f"[infer] ===== END t={time.perf_counter()-t0:.1f}s =====")