Refactor _infer_gpu into single-responsibility helper functions

Split the monolithic function into _log_infer_start, _log_gpu_properties,
_place_pipe_on_device, _instrument_first_touch, _make_step_callback, and
_log_infer_error. _infer_gpu itself is now just the orchestration of these
steps, with no behavior change.

app.py

@@ -387,72 +387,111 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
         raise
 
 
+def _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode):
+    print(f"[infer] ===== START =====")
+    print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode}")
+    print(f"[infer] prompt={repr(prompt[:120])}")
+
+
+def _log_gpu_properties(cuda_ok):
+    if not cuda_ok:
+        return None
+    p = torch.cuda.get_device_properties(0)
+    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()
+    return p
+
+
+# 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. On a generously-sized
+# allocation, move the full ~37GB pipeline to GPU in one shot (fast); but the
+# 47GB 2g.48gb MIG slice this Space actually gets measured a peak of 46.82GB
+# and still OOM'd on that path, wasting ~40s before falling back — so only
+# attempt it when there's real headroom above that, otherwise go straight to
+# enable_model_cpu_offload.
+_FAST_PATH_MIN_GB = 60
+
+
+def _place_pipe_on_device(cuda_ok, gpu_props, t0):
+    if getattr(pipe.transformer, "_hf_hook", None) is not None:
+        return  # already placed by an earlier call sharing this worker
+    if cuda_ok and gpu_props.total_memory / 1024**3 >= _FAST_PATH_MIN_GB:
+        try:
+            pipe.to(device)
+            print(f"[infer] moved full pipe to {device} — t={time.perf_counter()-t0:.1f}s")
+            return
+        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)")
+            return
+    pipe.enable_model_cpu_offload(device=device)
+    print(f"[infer] enabled cpu offload on {device} (slice too small for fast path)")
+
+
+def _instrument_first_touch(modules_with_names, t0):
+    """Install self-removing forward-pre-hooks that log the moment each module is first entered."""
+    def _make_hook(name, handle_box):
+        def _hook(mod, inputs):
+            print(f"[infer] first call into {name} — {_gpu_mem_str(True, sync=True)} | t={time.perf_counter()-t0:.1f}s")
+            handle_box["h"].remove()
+        return _hook
+    for module, name in modules_with_names:
+        handle_box = {}
+        handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
+
+
+def _make_step_callback(steps, timer, t0):
+    """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
+    step_times = []
+    def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
+        now = time.perf_counter()
+        step_times.append(now)
+        if step_idx == 0:
+            timer.mark("first_step")
+        timer.mark("last_step")  # overwritten each step; final value = end of last step
+        delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
+        tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
+        print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
+        return cb_kwargs
+    return _step_cb
+
+
+def _log_infer_error(e, t0, timer):
+    print(f"[infer] ERROR: {type(e).__name__}: {e} | t={time.perf_counter()-t0:.1f}s")
+    print(traceback.format_exc())
+    try:
+        torch.cuda.synchronize()
+    except Exception as cuda_err:
+        print(f"[infer] CUDA synchronize after error: {cuda_err}")
+    timer.print_timings()
+
+
 @spaces.GPU(duration=lambda *a, **kw: int(a[8]) if len(a) > 8 else 60)
 def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, gpu_duration=20):
     _cuda_ok = torch.cuda.is_available()
     timer = _InferTimer(_cuda_ok)
     t0 = time.perf_counter()
 
-    print(f"[infer] ===== START =====")
-    print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode}")
-    print(f"[infer] prompt={repr(prompt[:120])}")
+    _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode)
+    gpu_props = _log_gpu_properties(_cuda_ok)
 
-    if _cuda_ok:
-        p = torch.cuda.get_device_properties(0)
-        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()
-
-    # 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. On a
-    # generously-sized allocation, move the full ~37GB pipeline to GPU in one
-    # shot (fast); but the 47GB 2g.48gb MIG slice this Space actually gets
-    # measured a peak of 46.82GB and still OOM'd on that path, wasting ~40s
-    # before falling back — so only attempt it when there's real headroom
-    # above that, otherwise go straight to enable_model_cpu_offload.
-    _FAST_PATH_MIN_GB = 60
-    if getattr(pipe.transformer, "_hf_hook", None) is None:
-        if _cuda_ok and p.total_memory / 1024**3 >= _FAST_PATH_MIN_GB:
-            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)")
-        else:
-            pipe.enable_model_cpu_offload(device=device)
-            print(f"[infer] enabled cpu offload on {device} (slice too small for fast path)")
+    _place_pipe_on_device(_cuda_ok, gpu_props, t0)
     print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
 
     if _cuda_ok:
-        def _instrument_first_call(module, name):
-            handle_box = {}
-            def _hook(mod, inputs):
-                print(f"[infer] first call into {name} — {_gpu_mem_str(_cuda_ok, sync=True)} | t={time.perf_counter()-t0:.1f}s")
-                handle_box["h"].remove()
-            handle_box["h"] = module.register_forward_pre_hook(_hook)
-        _instrument_first_call(pipe.text_encoder, "text_encoder")
-        _instrument_first_call(pipe.transformer, "transformer")
-        _instrument_first_call(pipe.vae, "vae")
+        _instrument_first_touch(
+            [(pipe.text_encoder, "text_encoder"), (pipe.transformer, "transformer"), (pipe.vae, "vae")],
+            t0,
+        )
 
     print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
 
     generator = torch.Generator(device=device).manual_seed(seed)
-
-    _step_t = []
-    def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
-        now = time.perf_counter()
-        _step_t.append(now)
-        if step_idx == 0:
-            timer.mark("first_step")
-        timer.mark("last_step")  # overwritten each step; final value = end of last step
-        delta_ms = (now - (_step_t[-2] if len(_step_t) > 1 else t0)) * 1000
-        tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
-        print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
-        return cb_kwargs
+    step_cb = _make_step_callback(steps, timer, t0)
 
     timer.mark("pipe_start")
     print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")
@@ -461,7 +500,7 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
             image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
             height=height, width=width, num_inference_steps=steps,
             generator=generator, true_cfg_scale=guidance_scale,
-            callback_on_step_end=_step_cb,
+            callback_on_step_end=step_cb,
             callback_on_step_end_tensor_inputs=["latents"],
         ).images[0]
         timer.mark("pipe_end")
@@ -470,13 +509,7 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
         duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
         return result_image, seed, duration
     except Exception as e:
-        print(f"[infer] ERROR: {type(e).__name__}: {e} | t={time.perf_counter()-t0:.1f}s")
-        print(traceback.format_exc())
-        try:
-            torch.cuda.synchronize()
-        except Exception as cuda_err:
-            print(f"[infer] CUDA synchronize after error: {cuda_err}")
-        timer.print_timings()
+        _log_infer_error(e, t0, timer)
         raise
     finally:
         # No manual pipe.to("cpu"): in the offload-fallback case that fights the