Introduce a two-stage develop/main deploy flow (#27)

* Set up develop/main two-stage deploy: retarget dev-en to develop, add prod-en workflow

develop -> dev-en Space (staging), main -> prod-en Space (release).
See #26 for the release flow design.

* Centralize print_*() logging functions in logging_utils.py (#29)

Consolidates the named print_*() functions that #23 extracted in
model_loading.py, inference.py, examples_ui.py, image_codec.py, and the
qwenimage attention/pipeline modules into logging_utils.py, alongside the
print_log_*() functions already there, so all logging lives in one file
instead of being scattered per-module.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Rename dev-en's repo variable from HF_SPACE_ID to HF_SPACE_ID_DEV_EN (#30)

Matches the <ENV>_<LANG> naming already used by HF_SPACE_ID_PROD_EN and
HF_SPACE_ID_PROD_JA — the bare HF_SPACE_ID for dev-en was the odd one out.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Extract _InferTimer into a generic Timer in timing.py (#32)

print_timings() hardcoded a rows table of inference.py-specific mark
names, coupling an otherwise-reusable checkpoint timer to one call
site. Timer.print_report() now takes the row schema (and optional
total pair) as arguments instead of owning it.

Also drops the dead "image_load"/load_start/load_end row: nothing
ever marked those checkpoints, so that row and the overhead/total
lines never printed. The total now uses pipe_start/pipe_end, which
are actually marked.

Fixes #31

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

.github/workflows/deploy-dev-en.yml

@@ -2,7 +2,7 @@ name: Deploy to dev-en Space
 
 on:
   push:
-    branches: [main]
+    branches: [develop]
   workflow_dispatch: {}
 
 concurrency:
@@ -25,14 +25,14 @@ jobs:
       - name: Push to Hugging Face Space (dev-en)
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-          HF_SPACE_ID: ${{ vars.HF_SPACE_ID }}
+          HF_SPACE_ID: ${{ vars.HF_SPACE_ID_DEV_EN }}
         run: |
           if [ -z "$HF_TOKEN" ]; then
             echo "::error::HF_TOKEN secret is not set on this repo (Settings > Secrets and variables > Actions)."
             exit 1
           fi
           if [ -z "$HF_SPACE_ID" ]; then
-            echo "::error::HF_SPACE_ID variable is not set on this repo (Settings > Secrets and variables > Actions > Variables), format: <owner>/<space-name>."
+            echo "::error::HF_SPACE_ID_DEV_EN variable is not set on this repo (Settings > Secrets and variables > Actions > Variables), format: <owner>/<space-name>."
             exit 1
           fi
           git remote add space "https://user:${HF_TOKEN}@huggingface.co/spaces/${HF_SPACE_ID}"

.github/workflows/deploy-prod-en.yml

@@ -0,0 +1,42 @@
+name: Deploy to prod-en Space
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch: {}
+
+concurrency:
+  group: deploy-prod-en
+  cancel-in-progress: false
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+
+      - name: Push to Hugging Face Space (prod-en)
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          HF_SPACE_ID: ${{ vars.HF_SPACE_ID_PROD_EN }}
+        run: |
+          if [ -z "$HF_TOKEN" ]; then
+            echo "::error::HF_TOKEN secret is not set on this repo (Settings > Secrets and variables > Actions)."
+            exit 1
+          fi
+          if [ -z "$HF_SPACE_ID" ]; then
+            echo "::error::HF_SPACE_ID_PROD_EN variable is not set on this repo (Settings > Secrets and variables > Actions > Variables), format: <owner>/<space-name>."
+            exit 1
+          fi
+          git remote add space "https://user:${HF_TOKEN}@huggingface.co/spaces/${HF_SPACE_ID}"
+          # No --force: a non-fast-forward push means the Space diverged from GitHub
+          # (e.g. edited directly in the HF UI) and needs a human to reconcile it,
+          # rather than silently overwriting that history.
+          git push space HEAD:main

.github/workflows/typecheck.yml

@@ -2,7 +2,7 @@ name: Type check
 
 on:
   push:
-    branches: [main]
+    branches: [main, develop]
   pull_request: {}
 
 permissions:

examples_ui.py

@@ -6,6 +6,7 @@ import os
 from typing import Any
 
 from image_codec import encode_full_image, make_thumb_b64
+from logging_utils import print_built_example_cards, print_built_suggestion_chips, print_building_example_thumbnails
 
 with open("examples.json") as _f:
     EXAMPLES_CONFIG: list[dict[str, Any]] = json.load(_f)
@@ -73,18 +74,6 @@ def load_example_data(idx_str: str) -> str:
     return json.dumps({"images": b64_list, "prompt": ex["prompt"], "names": names, "status": "ok"})
 
 
-def print_building_example_thumbnails() -> None:
-    print("Building example thumbnails...")
-
-
-def print_built_example_cards(n: int) -> None:
-    print(f"Built {n} example cards.")
-
-
-def print_built_suggestion_chips(n: int) -> None:
-    print(f"Built {n} suggestion chips.")
-
-
 print_building_example_thumbnails()
 EXAMPLE_CARDS_HTML = build_example_cards_html()
 print_built_example_cards(len(EXAMPLES_CONFIG))

image_codec.py

@@ -7,19 +7,9 @@ from io import BytesIO
 from PIL import Image
 from PIL.Image import Image as PILImage
 
-LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
-
-
-def print_thumbnail_error(path: str, e: Exception) -> None:
-    print(f"Thumbnail error for {path}: {e}")
-
+from logging_utils import print_decode_error, print_encode_error, print_thumbnail_error
 
-def print_encode_error(path: str, e: Exception) -> None:
-    print(f"Encode error for {path}: {e}")
-
-
-def print_decode_error(e: Exception) -> None:
-    print(f"Error decoding image: {e}")
+LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
 
 
 def make_thumb_b64(path: str, max_dim: int = 220) -> str:

inference.py

@@ -18,9 +18,36 @@ from transformers import Qwen2_5_VLForConditionalGeneration
 
 from dimensions import compute_output_dimensions
 from image_codec import b64_to_pil_list
-from logging_utils import LogUploader
+from logging_utils import (
+    LogUploader,
+    print_bf16_text_encoder_moved,
+    print_calling_pipe,
+    print_cuda_sync_after_error,
+    print_first_call_into_module,
+    print_gpu_mem_status,
+    print_images_predecoded,
+    print_infer_end,
+    print_infer_error,
+    print_infer_exception,
+    print_infer_gpu_properties,
+    print_infer_params,
+    print_infer_prompt,
+    print_infer_start_header,
+    print_infer_traceback,
+    print_int8_text_encoder_load_failed,
+    print_int8_text_encoder_loaded,
+    print_pre_int8_load,
+    print_pre_vae_decode,
+    print_step_done,
+    print_text_encoder_offload_skipped_int8,
+    print_text_encoder_offload_skipped_mode,
+    print_text_encoder_offloaded,
+    print_vae_decode_done,
+    print_vae_tiling_activation,
+)
 from mode import Mode
 from model_loading import device, pipe
+from timing import Timer
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -35,65 +62,13 @@ _log_uploader = LogUploader(
 )
 
 
-def print_timing_divider() -> None:
-    print("[timing] ─────────────────────────────────────")
-
-
-def print_timing_lines(lines: list[str]) -> None:
-    print("\n".join(lines))
-
-
-class _InferTimer:
-    def __init__(self, cuda_ok: bool) -> None:
-        self._cuda_ok = cuda_ok
-        self._marks: dict[str, tuple[torch.cuda.Event | None, float]] = {}
-
-    def mark(self, name: str) -> None:
-        ev: torch.cuda.Event | None = None
-        if self._cuda_ok:
-            ev = torch.cuda.Event(enable_timing=True)  # type: ignore[no-untyped-call]
-            ev.record()
-        self._marks[name] = (ev, time.perf_counter())
-
-    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 float(ev_a.elapsed_time(ev_b))  # true GPU-timeline ms
-        return (t_b - t_a) * 1000.0
-
-    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"),
-        ]
-        total_ms = 0.0
-        lines = []
-        for label, a, b in rows:
-            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 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_divider()
-        print_timing_lines(lines)
-        print_timing_divider()
+# (label, start_mark, end_mark) rows for Timer.print_report(), matching the marks
+# _infer_gpu/_make_step_callback record: pipe_start, first_step, last_step, pipe_end.
+_TIMING_ROWS: list[tuple[str, str, str]] = [
+    ("preprocess", "pipe_start", "first_step"),
+    ("inference",  "first_step", "last_step"),
+    ("vae_decode", "last_step",  "pipe_end"),
+]
 
 
 def _gpu_mem_str(cuda_ok: bool, sync: bool = False) -> str:
@@ -169,26 +144,6 @@ def infer(images_b64_json: str, prompt: str, seed: int, randomize_seed: bool, gu
         raise
 
 
-def print_infer_exception(e: Exception) -> None:
-    print(f"[infer] EXCEPTION type={type(e).__module__}.{type(e).__qualname__} repr={e!r}")
-
-
-def print_infer_start_header() -> None:
-    print("[infer] ===== START =====")
-
-
-def print_infer_params(steps: int, guidance_scale: float, seed: int, gpu_duration: int, mode: Mode) -> None:
-    print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode.value}")
-
-
-def print_infer_prompt(prompt: str) -> None:
-    print(f"[infer] prompt={repr(prompt[:120])}")
-
-
-def print_infer_gpu_properties(p: torch.cuda._CudaDeviceProperties) -> None:
-    print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
-
-
 def _log_infer_start(prompt: str, steps: int, guidance_scale: float, seed: int, gpu_duration: int, mode: Mode) -> None:
     print_infer_start_header()
     print_infer_params(steps, guidance_scale, seed, gpu_duration, mode)
@@ -210,22 +165,6 @@ def _log_gpu_properties(cuda_ok: bool) -> torch.cuda._CudaDeviceProperties | Non
 _TEXT_ENCODER_INT8_REPO = os.environ.get("TEXT_ENCODER_INT8_REPO")
 
 
-def print_pre_int8_load(mem_str: str, elapsed: float) -> None:
-    print(f"[infer] pre-int8-load — {mem_str} | t={elapsed:.1f}s")
-
-
-def print_int8_text_encoder_loaded(repo: str, load_ms: float, elapsed: float, mem_str: str) -> None:
-    print(
-        f"[infer] loaded int8 text_encoder from {repo} — "
-        f"{load_ms:.0f}ms | t={elapsed:.1f}s | "
-        f"{mem_str}"
-    )
-
-
-def print_int8_text_encoder_load_failed(e: Exception) -> None:
-    print(f"[infer] WARNING: int8 text_encoder load failed, keeping bf16: {type(e).__name__}: {e}")
-
-
 def _ensure_int8_text_encoder(cuda_ok: bool, t0: float) -> None:
     # bitsandbytes 8bit modules can't be moved between devices with `.to()` (transformers
     # raises unconditionally for 8bit, unlike the version-gated allowance for 4bit), so this
@@ -252,10 +191,6 @@ def _ensure_int8_text_encoder(cuda_ok: bool, t0: float) -> None:
         print_int8_text_encoder_load_failed(e)
 
 
-def print_bf16_text_encoder_moved(device: torch.device, move_ms: float, elapsed: float) -> None:
-    print(f"[infer] moved bf16 text_encoder to {device} — {move_ms:.0f}ms | t={elapsed:.1f}s")
-
-
 def _ensure_bf16_text_encoder_on_device(cuda_ok: bool, t0: float) -> None:
     # Fallback for when int8 quantization is unavailable/unset/failed: model_loading.py
     # deliberately leaves the bf16 text_encoder on CPU at startup (see the comment in
@@ -271,10 +206,6 @@ def _ensure_bf16_text_encoder_on_device(cuda_ok: bool, t0: float) -> None:
     print_bf16_text_encoder_moved(device, (time.perf_counter() - _t0) * 1000, time.perf_counter() - t0)
 
 
-def print_first_call_into_module(name: str, mem_str: str, elapsed: float) -> None:
-    print(f"[infer] first call into {name} — {mem_str} | t={elapsed:.1f}s")
-
-
 def _instrument_first_touch(modules_with_names: list[tuple[torch.nn.Module, str]], t0: float) -> None:
     """Install self-removing forward-pre-hooks that log the moment each module is first entered."""
     def _make_hook(name: str, handle_box: dict[str, torch.utils.hooks.RemovableHandle]) -> Any:
@@ -287,27 +218,7 @@ def _instrument_first_touch(modules_with_names: list[tuple[torch.nn.Module, str]
         handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
 
 
-def print_step_done(step_idx: int, steps: int, delta_ms: float, tag: str, elapsed: float) -> None:
-    print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={elapsed:.1f}s")
-
-
-def print_text_encoder_offloaded(offload_ms: float, elapsed: float) -> None:
-    print(f"[infer] text_encoder offload to cpu — {offload_ms:.0f}ms | t={elapsed:.1f}s")
-
-
-def print_text_encoder_offload_skipped_int8() -> None:
-    print("[infer] skipping text_encoder offload (int8, .to() unsupported / smaller footprint)")
-
-
-def print_text_encoder_offload_skipped_mode(mode: Mode) -> None:
-    print(f"[infer] skipping text_encoder offload for mode={mode.value} (ample headroom at this resolution)")
-
-
-def print_pre_vae_decode(mem_str: str, elapsed: float) -> None:
-    print(f"[infer] pre-VAE-decode — {mem_str} | t={elapsed:.1f}s")
-
-
-def _make_step_callback(steps: int, timer: _InferTimer, t0: float, mode: Mode, cuda_ok: bool = False) -> Any:
+def _make_step_callback(steps: int, timer: Timer, t0: float, mode: Mode, cuda_ok: bool = False) -> Any:
     """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
     step_times: list[float] = []
     def _step_cb(pipeline: Any, step_idx: int, timestep: Any, cb_kwargs: dict[str, Any]) -> dict[str, Any]:
@@ -344,26 +255,14 @@ def _make_step_callback(steps: int, timer: _InferTimer, t0: float, mode: Mode, c
     return _step_cb
 
 
-def print_infer_error(e: Exception, elapsed: float) -> None:
-    print(f"[infer] ERROR: {type(e).__name__}: {e} | t={elapsed:.1f}s")
-
-
-def print_infer_traceback() -> None:
-    print(traceback.format_exc())
-
-
-def print_cuda_sync_after_error(cuda_err: Exception) -> None:
-    print(f"[infer] CUDA synchronize after error: {cuda_err}")
-
-
-def _log_infer_error(e: Exception, t0: float, timer: _InferTimer) -> None:
+def _log_infer_error(e: Exception, t0: float, timer: Timer) -> None:
     print_infer_error(e, time.perf_counter() - t0)
     print_infer_traceback()
     try:
         torch.cuda.synchronize()
     except Exception as cuda_err:
         print_cuda_sync_after_error(cuda_err)
-    timer.print_timings()
+    timer.print_report(_TIMING_ROWS, total=("pipe_start", "pipe_end"))
 
 
 # Every @spaces.GPU call lands on a fresh worker, so _ensure_int8_text_encoder's reload is
@@ -379,36 +278,11 @@ _COLD_START_BUFFER_S = 45
 _MAX_GPU_DURATION_S = 120  # matches the gpu_duration slider's max in the UI
 
 
-def print_gpu_mem_status(mem_str: str, elapsed: float) -> None:
-    print(f"[infer] {mem_str} — t={elapsed:.1f}s")
-
-
-def print_images_predecoded(n: int, width: int, height: int, seed: int) -> None:
-    print(f"[infer] {n} image(s) pre-decoded, output={width}x{height}, seed={seed}")
-
-
-def print_vae_tiling_activation(will_tile: bool, height: int, width: int) -> None:
-    print(f"[infer] VAE tiling will {'activate' if will_tile else 'NOT activate'} "
-          f"(threshold={height}x{width}px)")
-
-
-def print_calling_pipe(elapsed: float) -> None:
-    print(f"[infer] calling pipe... t={elapsed:.1f}s")
-
-
-def print_vae_decode_done(mem_str: str, elapsed: float) -> None:
-    print(f"[infer] VAE decode + postprocess done — {mem_str} | t={elapsed:.1f}s")
-
-
-def print_infer_end(elapsed: float) -> None:
-    print(f"[infer] ===== END t={elapsed:.1f}s =====")
-
-
 @spaces.GPU(duration=lambda *a, **kw: min(int(a[8]) + _COLD_START_BUFFER_S, _MAX_GPU_DURATION_S) if len(a) > 8 else 60)  # type: ignore[untyped-decorator]
 def _infer_gpu(pil_images: list[PILImage], prompt: str, seed: int, guidance_scale: float, steps: int,
                width: int, height: int, mode: Mode, gpu_duration: int = 20) -> tuple[PILImage, int, float]:
     _cuda_ok = torch.cuda.is_available()
-    timer = _InferTimer(_cuda_ok)
+    timer = Timer(_cuda_ok)
     t0 = time.perf_counter()
 
     _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode)
@@ -446,7 +320,7 @@ def _infer_gpu(pil_images: list[PILImage], prompt: str, seed: int, guidance_scal
         ).images[0]
         timer.mark("pipe_end")
         print_vae_decode_done(_gpu_mem_str(_cuda_ok, sync=True), time.perf_counter() - t0)
-        timer.print_timings()
+        timer.print_report(_TIMING_ROWS, total=("pipe_start", "pipe_end"))
         duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
         return result_image, seed, duration
     except Exception as e:

logging_utils.py

@@ -5,13 +5,258 @@ import time as _time
 from io import BytesIO
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
+import torch
 from huggingface_hub import hf_hub_download, CommitOperationAdd, CommitOperationDelete
 from PIL.Image import Image as PILImage
 
+from mode import Mode
+
 if TYPE_CHECKING:
     from huggingface_hub import HfApi
 
 
+def print_cuda_visible_devices() -> None:
+    print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"), flush=True)
+
+
+def print_torch_version() -> None:
+    print("torch.__version__ =", torch.__version__, flush=True)
+
+
+def print_using_device(device: torch.device) -> None:
+    print("Using device:", device, flush=True)
+
+
+def print_cuda_device_count() -> None:
+    print(f"CUDA device_count={torch.cuda.device_count()}, is_available={torch.cuda.is_available()}", flush=True)
+
+
+def print_env_gpu(p: torch.cuda._CudaDeviceProperties) -> None:
+    print(f"[env] GPU: {p.name}, VRAM={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}", flush=True)
+
+
+def print_env_cuda_version() -> None:
+    print(f"[env] CUDA (torch build): {torch.version.cuda}", flush=True)
+
+
+def print_env_cudnn_version() -> None:
+    print(f"[env] cuDNN: {torch.backends.cudnn.version()}", flush=True)  # type: ignore[no-untyped-call]
+
+
+def print_env_package_version(pkg: str, version: str) -> None:
+    print(f"[env] {pkg}=={version}", flush=True)
+
+
+def print_env_package_version_unavailable(pkg: str, error: Exception) -> None:
+    print(f"[env] {pkg}==? ({error})", flush=True)
+
+
+def print_env_ram(total_gb: float, avail_gb: float) -> None:
+    print(f"[env] RAM: {total_gb:.0f}GB total, {avail_gb:.0f}GB available", flush=True)
+
+
+def print_env_ram_unavailable(error: Exception) -> None:
+    print(f"[env] RAM: unavailable ({error})", flush=True)
+
+
+def print_tf32_enabled() -> None:
+    print("[startup] TF32 enabled", flush=True)
+
+
+def print_heartbeat(label: str, elapsed: float) -> None:
+    print(f"[startup] {label} still loading... ({elapsed:.0f}s)", flush=True)
+
+
+def print_unpatched_fp8_param(name: str, pname: str, module_type: str) -> None:
+    print(
+        f"[startup] WARNING: unpatched fp8 parameter {name}.{pname} "
+        f"({module_type}) — will likely error at inference",
+        flush=True,
+    )
+
+
+def print_loading_transformer() -> None:
+    print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
+
+
+def print_transformer_loaded(elapsed: float) -> None:
+    print(f"[startup] transformer loaded in {elapsed:.1f}s", flush=True)
+
+
+def print_transformer_patched(n_patched: int) -> None:
+    print(f"[startup] patched {n_patched} fp8-resident nn.Linear/RMSNorm modules for just-in-time upcast", flush=True)
+
+
+def print_transformer_memory_footprint(gb: float) -> None:
+    print(f"[startup] transformer memory footprint: {gb:.2f}GB", flush=True)
+
+
+def print_transformer_memory_footprint_unavailable(error: Exception) -> None:
+    print(f"[startup] transformer memory footprint: unavailable ({error})", flush=True)
+
+
+def print_loading_pipeline() -> None:
+    print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
+
+
+def print_vae_tiling(height: int, width: int, use_tiling: bool) -> None:
+    print(f"[startup] VAE tiling: threshold={height}x{width}px use_tiling={use_tiling}", flush=True)
+
+
+def print_pipeline_loaded(elapsed: float) -> None:
+    print(f"[startup] pipeline loaded in {elapsed:.1f}s", flush=True)
+
+
+def print_setting_attn_processor() -> None:
+    print("[startup] setting cuDNN SDPA attention processor...", flush=True)
+
+
+def print_attn_processor_set() -> None:
+    print("[startup] cuDNN SDPA attention processor set.", flush=True)
+
+
+def print_timing_divider() -> None:
+    print("[timing] ─────────────────────────────────────")
+
+
+def print_timing_lines(lines: list[str]) -> None:
+    print("\n".join(lines))
+
+
+def print_infer_exception(e: Exception) -> None:
+    print(f"[infer] EXCEPTION type={type(e).__module__}.{type(e).__qualname__} repr={e!r}")
+
+
+def print_infer_start_header() -> None:
+    print("[infer] ===== START =====")
+
+
+def print_infer_params(steps: int, guidance_scale: float, seed: int, gpu_duration: int, mode: Mode) -> None:
+    print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode.value}")
+
+
+def print_infer_prompt(prompt: str) -> None:
+    print(f"[infer] prompt={repr(prompt[:120])}")
+
+
+def print_infer_gpu_properties(p: torch.cuda._CudaDeviceProperties) -> None:
+    print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
+
+
+def print_pre_int8_load(mem_str: str, elapsed: float) -> None:
+    print(f"[infer] pre-int8-load — {mem_str} | t={elapsed:.1f}s")
+
+
+def print_int8_text_encoder_loaded(repo: str, load_ms: float, elapsed: float, mem_str: str) -> None:
+    print(
+        f"[infer] loaded int8 text_encoder from {repo} — "
+        f"{load_ms:.0f}ms | t={elapsed:.1f}s | "
+        f"{mem_str}"
+    )
+
+
+def print_int8_text_encoder_load_failed(e: Exception) -> None:
+    print(f"[infer] WARNING: int8 text_encoder load failed, keeping bf16: {type(e).__name__}: {e}")
+
+
+def print_bf16_text_encoder_moved(device: torch.device, move_ms: float, elapsed: float) -> None:
+    print(f"[infer] moved bf16 text_encoder to {device} — {move_ms:.0f}ms | t={elapsed:.1f}s")
+
+
+def print_first_call_into_module(name: str, mem_str: str, elapsed: float) -> None:
+    print(f"[infer] first call into {name} — {mem_str} | t={elapsed:.1f}s")
+
+
+def print_step_done(step_idx: int, steps: int, delta_ms: float, tag: str, elapsed: float) -> None:
+    print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={elapsed:.1f}s")
+
+
+def print_text_encoder_offloaded(offload_ms: float, elapsed: float) -> None:
+    print(f"[infer] text_encoder offload to cpu — {offload_ms:.0f}ms | t={elapsed:.1f}s")
+
+
+def print_text_encoder_offload_skipped_int8() -> None:
+    print("[infer] skipping text_encoder offload (int8, .to() unsupported / smaller footprint)")
+
+
+def print_text_encoder_offload_skipped_mode(mode: Mode) -> None:
+    print(f"[infer] skipping text_encoder offload for mode={mode.value} (ample headroom at this resolution)")
+
+
+def print_pre_vae_decode(mem_str: str, elapsed: float) -> None:
+    print(f"[infer] pre-VAE-decode — {mem_str} | t={elapsed:.1f}s")
+
+
+def print_infer_error(e: Exception, elapsed: float) -> None:
+    print(f"[infer] ERROR: {type(e).__name__}: {e} | t={elapsed:.1f}s")
+
+
+def print_infer_traceback() -> None:
+    import traceback
+    print(traceback.format_exc())
+
+
+def print_cuda_sync_after_error(cuda_err: Exception) -> None:
+    print(f"[infer] CUDA synchronize after error: {cuda_err}")
+
+
+def print_gpu_mem_status(mem_str: str, elapsed: float) -> None:
+    print(f"[infer] {mem_str} — t={elapsed:.1f}s")
+
+
+def print_images_predecoded(n: int, width: int, height: int, seed: int) -> None:
+    print(f"[infer] {n} image(s) pre-decoded, output={width}x{height}, seed={seed}")
+
+
+def print_vae_tiling_activation(will_tile: bool, height: int, width: int) -> None:
+    print(f"[infer] VAE tiling will {'activate' if will_tile else 'NOT activate'} "
+          f"(threshold={height}x{width}px)")
+
+
+def print_calling_pipe(elapsed: float) -> None:
+    print(f"[infer] calling pipe... t={elapsed:.1f}s")
+
+
+def print_vae_decode_done(mem_str: str, elapsed: float) -> None:
+    print(f"[infer] VAE decode + postprocess done — {mem_str} | t={elapsed:.1f}s")
+
+
+def print_infer_end(elapsed: float) -> None:
+    print(f"[infer] ===== END t={elapsed:.1f}s =====")
+
+
+def print_building_example_thumbnails() -> None:
+    print("Building example thumbnails...")
+
+
+def print_built_example_cards(n: int) -> None:
+    print(f"Built {n} example cards.")
+
+
+def print_built_suggestion_chips(n: int) -> None:
+    print(f"Built {n} suggestion chips.")
+
+
+def print_thumbnail_error(path: str, e: Exception) -> None:
+    print(f"Thumbnail error for {path}: {e}")
+
+
+def print_encode_error(path: str, e: Exception) -> None:
+    print(f"Encode error for {path}: {e}")
+
+
+def print_decode_error(e: Exception) -> None:
+    print(f"Error decoding image: {e}")
+
+
+def print_cudnn_sdpa_fallback(e: RuntimeError) -> None:
+    print(f"[attn] cuDNN SDPA backend unavailable ({e}), falling back to default", flush=True)
+
+
+def print_pipe_phase_timing(label: str, delta_ms: float, elapsed_s: float) -> None:
+    print(f"[pipe] {label} — {delta_ms:.0f}ms | t={elapsed_s:.1f}s", flush=True)
+
+
 def _img_to_jpeg(img: PILImage | None, quality: int = 85) -> bytes | None:
     if img is None:
         return None

model_loading.py

@@ -15,6 +15,32 @@ import spaces  # noqa: F401  (must be imported before any CUDA-touching code bel
 import torch
 from diffusers.models.normalization import RMSNorm
 
+from logging_utils import (
+    print_attn_processor_set,
+    print_cuda_device_count,
+    print_cuda_visible_devices,
+    print_env_cuda_version,
+    print_env_cudnn_version,
+    print_env_gpu,
+    print_env_package_version,
+    print_env_package_version_unavailable,
+    print_env_ram,
+    print_env_ram_unavailable,
+    print_heartbeat,
+    print_loading_pipeline,
+    print_loading_transformer,
+    print_pipeline_loaded,
+    print_setting_attn_processor,
+    print_tf32_enabled,
+    print_torch_version,
+    print_transformer_loaded,
+    print_transformer_memory_footprint,
+    print_transformer_memory_footprint_unavailable,
+    print_transformer_patched,
+    print_unpatched_fp8_param,
+    print_using_device,
+    print_vae_tiling,
+)
 from mode import Mode
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
@@ -23,110 +49,9 @@ from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.bfloat16
 
-
-def print_cuda_visible_devices() -> None:
-    print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"), flush=True)
-
-
-def print_torch_version() -> None:
-    print("torch.__version__ =", torch.__version__, flush=True)
-
-
-def print_using_device() -> None:
-    print("Using device:", device, flush=True)
-
-
-def print_cuda_device_count() -> None:
-    print(f"CUDA device_count={torch.cuda.device_count()}, is_available={torch.cuda.is_available()}", flush=True)
-
-
-def print_env_gpu(p: torch.cuda._CudaDeviceProperties) -> None:
-    print(f"[env] GPU: {p.name}, VRAM={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}", flush=True)
-
-
-def print_env_cuda_version() -> None:
-    print(f"[env] CUDA (torch build): {torch.version.cuda}", flush=True)
-
-
-def print_env_cudnn_version() -> None:
-    print(f"[env] cuDNN: {torch.backends.cudnn.version()}", flush=True)  # type: ignore[no-untyped-call]
-
-
-def print_env_package_version(pkg: str, version: str) -> None:
-    print(f"[env] {pkg}=={version}", flush=True)
-
-
-def print_env_package_version_unavailable(pkg: str, error: Exception) -> None:
-    print(f"[env] {pkg}==? ({error})", flush=True)
-
-
-def print_env_ram(total_gb: float, avail_gb: float) -> None:
-    print(f"[env] RAM: {total_gb:.0f}GB total, {avail_gb:.0f}GB available", flush=True)
-
-
-def print_env_ram_unavailable(error: Exception) -> None:
-    print(f"[env] RAM: unavailable ({error})", flush=True)
-
-
-def print_tf32_enabled() -> None:
-    print("[startup] TF32 enabled", flush=True)
-
-
-def print_heartbeat(label: str, elapsed: float) -> None:
-    print(f"[startup] {label} still loading... ({elapsed:.0f}s)", flush=True)
-
-
-def print_unpatched_fp8_param(name: str, pname: str, module_type: str) -> None:
-    print(
-        f"[startup] WARNING: unpatched fp8 parameter {name}.{pname} "
-        f"({module_type}) — will likely error at inference",
-        flush=True,
-    )
-
-
-def print_loading_transformer() -> None:
-    print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
-
-
-def print_transformer_loaded(elapsed: float) -> None:
-    print(f"[startup] transformer loaded in {elapsed:.1f}s", flush=True)
-
-
-def print_transformer_patched(n_patched: int) -> None:
-    print(f"[startup] patched {n_patched} fp8-resident nn.Linear/RMSNorm modules for just-in-time upcast", flush=True)
-
-
-def print_transformer_memory_footprint(gb: float) -> None:
-    print(f"[startup] transformer memory footprint: {gb:.2f}GB", flush=True)
-
-
-def print_transformer_memory_footprint_unavailable(error: Exception) -> None:
-    print(f"[startup] transformer memory footprint: unavailable ({error})", flush=True)
-
-
-def print_loading_pipeline() -> None:
-    print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
-
-
-def print_vae_tiling(height: int, width: int, use_tiling: bool) -> None:
-    print(f"[startup] VAE tiling: threshold={height}x{width}px use_tiling={use_tiling}", flush=True)
-
-
-def print_pipeline_loaded(elapsed: float) -> None:
-    print(f"[startup] pipeline loaded in {elapsed:.1f}s", flush=True)
-
-
-def print_setting_attn_processor() -> None:
-    print("[startup] setting cuDNN SDPA attention processor...", flush=True)
-
-
-def print_attn_processor_set() -> None:
-    print("[startup] cuDNN SDPA attention processor set.", flush=True)
-
-
 print_cuda_visible_devices()
 print_torch_version()
-print_using_device()
+print_using_device(device)
 print_cuda_device_count()
 
 

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -30,6 +30,8 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
 
+from logging_utils import print_pipe_phase_timing
+
 
 if is_torch_xla_available():
     import torch_xla.core.xla_model as xm
@@ -156,10 +158,6 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-def print_pipe_phase_timing(label, delta_ms, elapsed_s):
-    print(f"[pipe] {label} — {delta_ms:.0f}ms | t={elapsed_s:.1f}s", flush=True)
-
-
 def calculate_dimensions(target_area, ratio):
     width = math.sqrt(target_area * ratio)
     height = width / ratio

qwenimage/qwen_fa3_processor.py

@@ -4,9 +4,7 @@ from typing import Optional, Tuple
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
 
-
-def print_cudnn_sdpa_fallback(e: RuntimeError) -> None:
-    print(f"[attn] cuDNN SDPA backend unavailable ({e}), falling back to default", flush=True)
+from logging_utils import print_cudnn_sdpa_fallback
 
 
 class QwenDoubleStreamAttnProcessorFA3:

timing.py

@@ -0,0 +1,64 @@
+"""Generic checkpoint timer for instrumenting wall-clock/GPU-timeline durations."""
+import time
+
+import torch
+
+from logging_utils import print_timing_divider, print_timing_lines
+
+
+class Timer:
+    """Records named marks and reports elapsed time between them.
+
+    On CUDA, each mark also captures a CUDA event so elapsed_ms() reports true
+    GPU-timeline duration instead of wall-clock time that includes async queuing.
+    """
+
+    def __init__(self, cuda_ok: bool) -> None:
+        self._cuda_ok = cuda_ok
+        self._marks: dict[str, tuple[torch.cuda.Event | None, float]] = {}
+
+    def mark(self, name: str) -> None:
+        ev: torch.cuda.Event | None = None
+        if self._cuda_ok:
+            ev = torch.cuda.Event(enable_timing=True)  # type: ignore[no-untyped-call]
+            ev.record()
+        self._marks[name] = (ev, time.perf_counter())
+
+    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 float(ev_a.elapsed_time(ev_b))  # true GPU-timeline ms
+        return (t_b - t_a) * 1000.0
+
+    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_report(self, rows: list[tuple[str, str, str]], total: tuple[str, str] | None = None) -> None:
+        """Print elapsed_ms() for each (label, start_mark, end_mark) row that has both marks recorded.
+
+        If `total` is given as (start_mark, end_mark) and both are recorded, also prints an
+        "overhead" line (total minus the sum of the printed rows) and a grand total line.
+        """
+        if self._cuda_ok:
+            try:
+                torch.cuda.synchronize()
+            except Exception:
+                pass
+        total_ms = 0.0
+        lines = []
+        for label, a, b in rows:
+            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 total and total[0] in self._marks and total[1] in self._marks:
+            overall_ms = self.elapsed_ms(*total)
+            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_divider()
+        print_timing_lines(lines)
+        print_timing_divider()