Refactor app.py: extract helpers, consolidate imports

- _InferTimer class, _gpu_mem_str, _validate_infer_inputs, _resolve_seed,
_spawn_log all extracted from infer()
- _example_thumbs_html, _example_card_html extracted from build_example_cards_html()
- _parse_example_idx extracted from load_example_data()
- Consolidate scattered imports (time, traceback) to module top
- Move negative_prompt load above infer() where it is used

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

app.py

@@ -1,6 +1,8 @@
 import os
 import gc
+import time
 import threading
+import traceback
 
 import gradio as gr
 import numpy as np
@@ -44,24 +46,23 @@ print("[startup] all imports done", flush=True)
 
 dtype = torch.bfloat16
 
-import time as _time
-_t0_load = _time.perf_counter()
+_t0_load = time.perf_counter()
 print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
 _transformer = QwenImageTransformer2DModel.from_pretrained(
     "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
     torch_dtype=dtype,
     device_map="cuda",
 )
-print(f"[startup] transformer loaded in {_time.perf_counter()-_t0_load:.1f}s", flush=True)
+print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", flush=True)
 
-_t1_load = _time.perf_counter()
+_t1_load = time.perf_counter()
 print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "FireRedTeam/FireRed-Image-Edit-1.1",
     transformer=_transformer,
     torch_dtype=dtype,
 ).to(device)
-print(f"[startup] pipeline loaded in {_time.perf_counter()-_t1_load:.1f}s", flush=True)
+print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
 
 print("Using default attention processor (FA3 skipped for ZeroGPU GPU-arch compatibility).", flush=True)
 
@@ -99,34 +100,44 @@ def encode_full_image(path):
         return ""
 
 
-def build_example_cards_html():
-    cards = ""
-    for i, ex in enumerate(EXAMPLES_CONFIG):
-        thumbs_html = ""
-        for path in ex["images"]:
-            thumb = make_thumb_b64(path)
-            if thumb:
-                thumbs_html += f'<img src="{thumb}" alt="">'
-            else:
-                thumbs_html += '<div class="example-thumb-placeholder">Preview</div>'
-        n = len(ex["images"])
-        badge = f'{n} image{"s" if n > 1 else ""}'
-        prompt_short = html_lib.escape(ex["prompt"][:90])
-        if len(ex["prompt"]) > 90:
-            prompt_short += "..."
-        cards += f'''<div class="example-card" data-idx="{i}">
+def _example_thumbs_html(images):
+    html = ""
+    for path in images:
+        thumb = make_thumb_b64(path)
+        if thumb:
+            html += f'<img src="{thumb}" alt="">'
+        else:
+            html += '<div class="example-thumb-placeholder">Preview</div>'
+    return html
+
+
+def _example_card_html(idx, ex):
+    thumbs_html = _example_thumbs_html(ex["images"])
+    n = len(ex["images"])
+    badge = f'{n} image{"s" if n > 1 else ""}'
+    prompt_short = html_lib.escape(ex["prompt"][:90])
+    if len(ex["prompt"]) > 90:
+        prompt_short += "..."
+    return f'''<div class="example-card" data-idx="{idx}">
             <div class="example-thumbs">{thumbs_html}</div>
             <div class="example-meta"><span class="example-badge">{badge}</span></div>
             <div class="example-prompt-text">{prompt_short}</div>
         </div>'''
-    return cards
 
 
-def load_example_data(idx_str):
+def build_example_cards_html():
+    return "".join(_example_card_html(i, ex) for i, ex in enumerate(EXAMPLES_CONFIG))
+
+
+def _parse_example_idx(idx_str):
     try:
-        idx = int(float(idx_str)) if idx_str and idx_str.strip() else -1
+        return int(float(idx_str)) if idx_str and idx_str.strip() else -1
     except (ValueError, TypeError):
-        idx = -1
+        return -1
+
+
+def load_example_data(idx_str):
+    idx = _parse_example_idx(idx_str)
     if idx < 0 or idx >= len(EXAMPLES_CONFIG):
         return json.dumps({"images": [], "prompt": "", "names": [], "status": "error"})
     ex = EXAMPLES_CONFIG[idx]
@@ -174,70 +185,123 @@ def update_dimensions_on_upload(image):
     return compute_output_dimensions(w, h)
 
 
-@spaces.GPU
-def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
-    import time, traceback
-    _cuda_ok = torch.cuda.is_available()
-    _marks: dict = {}  # name -> (cuda_event | None, perf_counter_float)
+class _InferTimer:
+    def __init__(self, cuda_ok: bool) -> None:
+        self._cuda_ok = cuda_ok
+        self._marks: dict = {}
 
-    def _mark(name: str) -> None:
+    def mark(self, name: str) -> None:
         ev = None
-        if _cuda_ok:
+        if self._cuda_ok:
             ev = torch.cuda.Event(enable_timing=True)
             ev.record()
-        _marks[name] = (ev, time.perf_counter())
+        self._marks[name] = (ev, time.perf_counter())
 
-    def _elapsed_ms(a: str, b: str) -> float:
-        ev_a, t_a = _marks[a]
-        ev_b, t_b = _marks[b]
+    def elapsed_ms(self, a: str, b: str) -> float:
+        ev_a, t_a = self._marks[a]
+        ev_b, t_b = self._marks[b]
         if ev_a and ev_b:
             return ev_a.elapsed_time(ev_b)  # true GPU-timeline ms
         return (t_b - t_a) * 1000.0
 
-    def _print_timings() -> None:
-        if _cuda_ok:
+    def wall_start(self, name: str) -> float:
+        return self._marks[name][1]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._marks
+
+    def print_timings(self) -> None:
+        if self._cuda_ok:
             try:
                 torch.cuda.synchronize()
             except Exception:
                 pass
         rows = [
-            ("image_load",  "load_start",  "load_end"),
-            ("preprocess",  "pipe_start",  "first_step"),
-            ("inference",   "first_step",  "last_step"),
-            ("vae_decode",  "last_step",   "pipe_end"),
+            ("image_load", "load_start", "load_end"),
+            ("preprocess", "pipe_start", "first_step"),
+            ("inference",  "first_step", "last_step"),
+            ("vae_decode", "last_step",  "pipe_end"),
         ]
         total_ms = 0.0
         lines = []
         for label, a, b in rows:
-            if a in _marks and b in _marks:
-                ms = _elapsed_ms(a, b)
+            if a in self._marks and b in self._marks:
+                ms = self.elapsed_ms(a, b)
                 total_ms += ms
                 lines.append(f"[timing]   {label:<14} {ms:8.1f} ms")
-        if "load_start" in _marks and "pipe_end" in _marks:
-            overall_ms = _elapsed_ms("load_start", "pipe_end")
+        if "load_start" in self._marks and "pipe_end" in self._marks:
+            overall_ms = self.elapsed_ms("load_start", "pipe_end")
             lines.append(f"[timing]   {'overhead':<14} {overall_ms - total_ms:8.1f} ms")
             lines.append(f"[timing]   {'── total ──':<14} {overall_ms:8.1f} ms")
         print("[timing] ─────────────────────────────────────")
         print("\n".join(lines))
         print("[timing] ─────────────────────────────────────")
 
-    t0 = time.perf_counter()
 
-    def _t():
-        return f"t={time.perf_counter()-t0:.1f}s"
+def _gpu_mem_str(cuda_ok: bool, sync: bool = False) -> str:
+    if not cuda_ok:
+        return "CUDA not available"
+    if sync:
+        try:
+            torch.cuda.synchronize()
+        except Exception as se:
+            return f"CUDA sync failed: {se}"
+    alloc    = torch.cuda.memory_allocated()    / 1024**3
+    reserved = torch.cuda.memory_reserved()     / 1024**3
+    peak     = torch.cuda.max_memory_allocated() / 1024**3
+    return f"alloc={alloc:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB"
 
-    def _mem(sync=False):
-        if not _cuda_ok:
-            return "CUDA not available"
-        if sync:
-            try:
-                torch.cuda.synchronize()
-            except Exception as se:
-                return f"CUDA sync failed: {se}"
-        alloc    = torch.cuda.memory_allocated()   / 1024**3
-        reserved = torch.cuda.memory_reserved()    / 1024**3
-        peak     = torch.cuda.max_memory_allocated()/ 1024**3
-        return f"alloc={alloc:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB"
+
+def _validate_infer_inputs(pil_images: list, prompt: str) -> None:
+    if not pil_images:
+        raise gr.Error("Please upload at least one image to edit.")
+    if not prompt or prompt.strip() == "":
+        raise gr.Error("Please enter an edit prompt.")
+
+
+def _resolve_seed(seed: int, randomize_seed: bool) -> int:
+    return random.randint(0, MAX_SEED) if randomize_seed else seed
+
+
+def _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale,
+               width, height, duration, success, error=""):
+    threading.Thread(
+        target=log_inference,
+        args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
+              width, height, duration, success, error),
+        daemon=True,
+    ).start()
+
+
+# ── static assets ─────────────────────────────────────────────────────────────
+
+with open("static/app.css") as _f:
+    css = _f.read()
+
+with open("static/gallery.js") as _f:
+    gallery_js = _f.read()
+
+with open("static/wire_outputs.js") as _f:
+    wire_outputs_js = _f.read()
+
+with open("static/run_preprocess.js") as _f:
+    run_preprocess_js = _f.read()
+
+with open("static/negative_prompt.txt") as _f:
+    negative_prompt = _f.read().strip()
+
+# ── HTML template ──────────────────────────────────────────────────────────────
+
+with open("templates/app.html") as _f:
+    app_html = _f.read().format(example_cards_html=EXAMPLE_CARDS_HTML)
+
+# ── Gradio blocks ──────────────────────────────────────────────────────────────
+
+@spaces.GPU
+def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+    _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}, randomize={randomize_seed}")
@@ -248,23 +312,20 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
         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()
 
-    print(f"[infer] {_mem()} — {_t()}")
+    print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
 
     gc.collect()
     torch.cuda.empty_cache()
-    print(f"[infer] cache cleared — {_mem()}")
+    print(f"[infer] cache cleared — {_gpu_mem_str(_cuda_ok)}")
 
-    _mark("load_start")
+    timer.mark("load_start")
     pil_images = b64_to_pil_list(images_b64_json)
-    _mark("load_end")
-    print(f"[infer] decoded {len(pil_images)} image(s) — {_elapsed_ms('load_start', 'load_end'):.0f}ms")
-    if not pil_images:
-        raise gr.Error("Please upload at least one image to edit.")
-    if not prompt or prompt.strip() == "":
-        raise gr.Error("Please enter an edit prompt.")
+    timer.mark("load_end")
+    print(f"[infer] decoded {len(pil_images)} image(s) — {timer.elapsed_ms('load_start', 'load_end'):.0f}ms")
+
+    _validate_infer_inputs(pil_images, prompt)
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    seed = _resolve_seed(seed, randomize_seed)
     generator = torch.Generator(device=device).manual_seed(seed)
     width, height = update_dimensions_on_upload(pil_images[0])
     print(f"[infer] input={pil_images[0].size}, output={width}x{height}, seed={seed}")
@@ -278,15 +339,15 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
         now = time.perf_counter()
         _step_t.append(now)
         if step_idx == 0:
-            _mark("first_step")
-        _mark("last_step")  # overwritten each iteration; final value = end of last step
+            timer.mark("first_step")
+        timer.mark("last_step")  # overwritten each iteration; final value = end of last step
         delta_ms = (now - (_step_t[-2] if len(_step_t) > 1 else t0)) * 1000
         tag = " ← includes compile" if step_idx == 0 else ""
-        print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | {_mem()} | {_t()}")
+        print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | {_gpu_mem_str(_cuda_ok)} | t={time.perf_counter()-t0:.1f}s")
         return cb_kwargs
 
-    _mark("pipe_start")
-    print(f"[infer] calling pipe... {_t()}")
+    timer.mark("pipe_start")
+    print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")
     try:
         result_image = pipe(
             image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
@@ -295,62 +356,28 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
             callback_on_step_end=_step_cb,
             callback_on_step_end_tensor_inputs=["latents"],
         ).images[0]
-        _mark("pipe_end")
-        print(f"[infer] VAE decode + postprocess done — {_mem(sync=True)} | {_t()}")
-        _print_timings()
-        duration = _elapsed_ms("pipe_start", "pipe_end") / 1000.0
-        threading.Thread(
-            target=log_inference,
-            args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
-                  width, height, duration, True, ""),
-            daemon=True,
-        ).start()
+        timer.mark("pipe_end")
+        print(f"[infer] VAE decode + postprocess done — {_gpu_mem_str(_cuda_ok, sync=True)} | t={time.perf_counter()-t0:.1f}s")
+        timer.print_timings()
+        duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
+        _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale, width, height, duration, True)
         return result_image, seed
     except Exception as e:
-        print(f"[infer] ERROR: {type(e).__name__}: {e} | {_t()}")
+        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}")
-        _print_timings()
-        duration = time.perf_counter() - _marks["pipe_start"][1]
-        threading.Thread(
-            target=log_inference,
-            args=(pil_images, None, prompt, seed, steps, guidance_scale,
-                  width, height, duration, False, str(e)),
-            daemon=True,
-        ).start()
+        timer.print_timings()
+        duration = time.perf_counter() - timer.wall_start("pipe_start")
+        _spawn_log(pil_images, None, prompt, seed, steps, guidance_scale, width, height, duration, False, str(e))
         raise e
     finally:
         gc.collect()
         torch.cuda.empty_cache()
-        print(f"[infer] ===== END {_t()} =====")
-
-
-# ── static assets ─────────────────────────────────────────��───────────────────
+        print(f"[infer] ===== END t={time.perf_counter()-t0:.1f}s =====")
 
-with open("static/app.css") as _f:
-    css = _f.read()
-
-with open("static/gallery.js") as _f:
-    gallery_js = _f.read()
-
-with open("static/wire_outputs.js") as _f:
-    wire_outputs_js = _f.read()
-
-with open("static/run_preprocess.js") as _f:
-    run_preprocess_js = _f.read()
-
-with open("static/negative_prompt.txt") as _f:
-    negative_prompt = _f.read().strip()
-
-# ── HTML template ──────────────────────────────────────────────────────────────
-
-with open("templates/app.html") as _f:
-    app_html = _f.read().format(example_cards_html=EXAMPLE_CARDS_HTML)
-
-# ── Gradio blocks ──────────────────────────────────────────────────────────────
 
 with gr.Blocks() as demo: