Skip text_encoder CPU offload in fast mode; type mode as an enum

Logged phase timing showed the text_encoder-to-CPU offload before VAE
decode costs ~9.7s, while decode itself takes ~113ms. That offload was
freeing memory decode never needed at fast mode's small resolution
(measured peak stayed well under the MIG slice's headroom) — it exists
to prevent OOM at high-detail's much larger 2048px decode, so it's kept
there and only skipped for fast mode.

Replaced the raw "fast"/"high_detail" string comparisons scattered
across dimensions.py and app.py with a Mode enum (mode.py) that owns
both behaviors (max_dim, offloads_text_encoder_before_decode) as
properties on the variant.

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

app.py

@@ -70,7 +70,8 @@ torch.backends.cudnn.allow_tf32 = True
 print("[startup] TF32 enabled", flush=True)
 
 print("[startup] importing dimensions...", flush=True)
-from dimensions import compute_output_dimensions, max_dim_for_mode, MAX_OUTPUT_DIM_FAST
+from dimensions import compute_output_dimensions
+from mode import Mode
 print("[startup] importing diffusers...", flush=True)
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.models.normalization import RMSNorm
@@ -189,7 +190,7 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     torch_dtype=dtype,
 )
 _hb.set()
-pipe.vae.enable_tiling(tile_sample_min_height=MAX_OUTPUT_DIM_FAST, tile_sample_min_width=MAX_OUTPUT_DIM_FAST)
+pipe.vae.enable_tiling(tile_sample_min_height=Mode.FAST.max_dim, tile_sample_min_width=Mode.FAST.max_dim)
 print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
 
 print("[startup] setting cuDNN SDPA attention processor...", flush=True)
@@ -448,10 +449,11 @@ with open("templates/app.html") as _f:
 def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps, mode, gpu_duration=20, progress=gr.Progress(track_tqdm=True)):
     # CPU-only preprocessing — GPU not yet allocated
     gc.collect()
+    mode = Mode.from_value(mode)
     pil_images = b64_to_pil_list(images_b64_json)
     _validate_infer_inputs(pil_images, prompt)
     seed = _resolve_seed(seed, randomize_seed)
-    width, height = update_dimensions_on_upload(pil_images[0], max_dim_for_mode(mode))
+    width, height = update_dimensions_on_upload(pil_images[0], mode.max_dim)
     t0 = time.perf_counter()
     try:
         result_image, seed, duration = _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, int(gpu_duration))
@@ -472,9 +474,9 @@ 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):
+def _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode: 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] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode.value}")
     print(f"[infer] prompt={repr(prompt[:120])}")
 
 
@@ -531,7 +533,7 @@ def _instrument_first_touch(modules_with_names, t0):
         handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
 
 
-def _make_step_callback(steps, timer, t0):
+def _make_step_callback(steps, timer, t0, mode: Mode):
     """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):
@@ -541,15 +543,19 @@ def _make_step_callback(steps, timer, t0):
             timer.mark("first_step")
         timer.mark("last_step")  # overwritten each step; final value = end of last step
         # Text encoder is done after prompt encoding, before the denoising loop starts.
-        # Drop it (~15GB) ahead of VAE decode's fp32-upcast memory spike. Skipped when
-        # accelerate hooks are managing placement (offload-fallback path) to avoid
-        # fighting their own device bookkeeping — see the finally block below.
+        # Dropping it (~15GB) ahead of VAE decode's fp32-upcast memory spike only pays off
+        # when that spike is big enough to need the headroom (see Mode.offloads_text_encoder_before_decode).
+        # Also skipped when accelerate hooks are managing placement (offload-fallback path)
+        # to avoid fighting their own device bookkeeping.
         if step_idx == steps - 1 and getattr(pipeline.text_encoder, "_hf_hook", None) is None:
-            _offload_t0 = time.perf_counter()
-            pipeline.text_encoder.to("cpu")
-            torch.cuda.empty_cache()
-            _offload_ms = (time.perf_counter() - _offload_t0) * 1000
-            print(f"[infer] text_encoder offload to cpu — {_offload_ms:.0f}ms | t={time.perf_counter()-t0:.1f}s")
+            if mode.offloads_text_encoder_before_decode:
+                _offload_t0 = time.perf_counter()
+                pipeline.text_encoder.to("cpu")
+                torch.cuda.empty_cache()
+                _offload_ms = (time.perf_counter() - _offload_t0) * 1000
+                print(f"[infer] text_encoder offload to cpu — {_offload_ms:.0f}ms | t={time.perf_counter()-t0:.1f}s")
+            else:
+                print(f"[infer] skipping text_encoder offload for mode={mode.value} (ample headroom at this resolution)")
         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")
@@ -568,7 +574,7 @@ def _log_infer_error(e, t0, timer):
 
 
 @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):
+def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode: Mode, gpu_duration=20):
     _cuda_ok = torch.cuda.is_available()
     timer = _InferTimer(_cuda_ok)
     t0 = time.perf_counter()
@@ -588,7 +594,7 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
     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_cb = _make_step_callback(steps, timer, t0)
+    step_cb = _make_step_callback(steps, timer, t0, mode)
 
     timer.mark("pipe_start")
     print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")

dimensions.py

@@ -2,9 +2,6 @@ MAX_OUTPUT_DIM = 2048
 MAX_OUTPUT_DIM_FAST = 768
 
 
-def max_dim_for_mode(mode): return MAX_OUTPUT_DIM_FAST if mode == "fast" else MAX_OUTPUT_DIM
-
-
 def compute_output_dimensions(w, h, max_dim=MAX_OUTPUT_DIM):
     # Pin the long side to max_dim and scale the short side proportionally.
     # We snap to the nearest multiple of 8 (not floor) to minimise aspect ratio

mode.py

@@ -0,0 +1,31 @@
+from enum import Enum
+
+from dimensions import MAX_OUTPUT_DIM, MAX_OUTPUT_DIM_FAST
+
+
+class Mode(str, Enum):
+    """UI speed/quality presets. Behavior that depends on the mode lives here as a
+    property on the variant, rather than call sites branching on a raw mode string."""
+
+    FAST = "fast"
+    HIGH_DETAIL = "high_detail"
+
+    @classmethod
+    def from_value(cls, value) -> "Mode":
+        try:
+            return cls(value)
+        except ValueError:
+            return cls.HIGH_DETAIL
+
+    @property
+    def max_dim(self) -> int:
+        return MAX_OUTPUT_DIM_FAST if self is Mode.FAST else MAX_OUTPUT_DIM
+
+    @property
+    def offloads_text_encoder_before_decode(self) -> bool:
+        # Fast mode's small decode (<=768px) fits comfortably in the headroom already
+        # freed by the fp8-resident transformer, so evicting the ~15GB text encoder
+        # first — measured at ~9.7s on this Space's MIG slice — buys nothing there.
+        # High-detail's decode is up to ~8.5x more pixels (2048px), where that
+        # headroom margin is unverified, so it keeps the safety net.
+        return self is Mode.HIGH_DETAIL

tests/test_dimensions.py

@@ -1,5 +1,5 @@
 import pytest
-from dimensions import compute_output_dimensions, max_dim_for_mode, MAX_OUTPUT_DIM, MAX_OUTPUT_DIM_FAST
+from dimensions import compute_output_dimensions, MAX_OUTPUT_DIM
 
 
 def aspect_ratio_error(w_in, h_in, w_out, h_out):
@@ -97,14 +97,3 @@ def test_custom_max_dim():
     nw, nh = compute_output_dimensions(1920, 1080, max_dim=1024)
     assert nw == 1024
     assert nh % 8 == 0
-
-
-# --- mode -> max dimension mapping ---
-
-def test_fast_mode_uses_fast_max_dim():
-    assert max_dim_for_mode("fast") == MAX_OUTPUT_DIM_FAST
-
-
-@pytest.mark.parametrize("mode", ["high_detail", "anything_else", None])
-def test_non_fast_modes_use_default_max_dim(mode):
-    assert max_dim_for_mode(mode) == MAX_OUTPUT_DIM

tests/test_mode.py

@@ -0,0 +1,40 @@
+import pytest
+from dimensions import MAX_OUTPUT_DIM, MAX_OUTPUT_DIM_FAST
+from mode import Mode
+
+
+# --- from_value ---
+
+def test_from_value_accepts_known_strings():
+    assert Mode.from_value("fast") is Mode.FAST
+    assert Mode.from_value("high_detail") is Mode.HIGH_DETAIL
+
+
+@pytest.mark.parametrize("value", ["anything_else", None, "", "FAST"])
+def test_from_value_falls_back_to_high_detail(value):
+    assert Mode.from_value(value) is Mode.HIGH_DETAIL
+
+
+def test_from_value_is_idempotent_on_mode_instances():
+    assert Mode.from_value(Mode.FAST) is Mode.FAST
+    assert Mode.from_value(Mode.HIGH_DETAIL) is Mode.HIGH_DETAIL
+
+
+# --- max_dim ---
+
+def test_fast_mode_uses_fast_max_dim():
+    assert Mode.FAST.max_dim == MAX_OUTPUT_DIM_FAST
+
+
+def test_high_detail_mode_uses_default_max_dim():
+    assert Mode.HIGH_DETAIL.max_dim == MAX_OUTPUT_DIM
+
+
+# --- offloads_text_encoder_before_decode ---
+
+def test_fast_mode_skips_text_encoder_offload():
+    assert Mode.FAST.offloads_text_encoder_before_decode is False
+
+
+def test_high_detail_mode_keeps_text_encoder_offload():
+    assert Mode.HIGH_DETAIL.offloads_text_encoder_before_decode is True