Keep transformer weights fp8-resident instead of upcasting at load

The Rapid-AIO-V23 checkpoint already ships natively in fp8_e4m3fn
(~19GB for 20B params), but from_pretrained's torch_dtype=bf16 was
casting every tensor up to bf16 at load time (~38GB), which is why the
pipeline (~35-40GB transformer + ~14GB text encoder) never comfortably
fit this Space's 47GB MIG slice and needed the slow enable_model_cpu_offload
fallback (_FAST_PATH_MIN_GB=60, gating out the 47GB slice entirely).

This keeps torch_dtype=torch.float8_e4m3fn on the transformer load (a
diffusers-supported path — see modeling_utils.py's explicit float8
handling) so weights stay fp8 in memory, and patches every fp8-resident
nn.Linear to upcast its own weight to the input's dtype just before the
matmul. nn.Linear has no fp8 GEMM kernel on this GPU, so this is not a
compute speedup — it's the same bf16 upcast the old code did, just
deferred from load-time (all layers at once) to call-time (one layer's
weight transiently bf16 at a time), so resident GPU memory drops from
~38GB to ~19GB for the transformer with no change to the actual math or
output. That should bring full-pipeline residency down to ~35GB,
comfortably under the 47GB slice, so _FAST_PATH_MIN_GB is lowered from
60 to 40 to let this Space take the fast pipe.to(device) path instead
of falling back to cpu offload. The OOM->offload fallback stays in
place as a safety net in case the estimate is off.

Needs real-GPU verification on the Space before merging to main — no
local GPU available to test weight-loading/placement changes.

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

app.py

@@ -3,6 +3,7 @@ import gc
 import time
 import threading
 import traceback
+import types
 
 # cudaMallocAsync bypasses NVML memory queries that fail on MIG GPU instances
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
@@ -94,16 +95,46 @@ def _start_heartbeat(label: str) -> threading.Event:
     return done
 
 
+_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
+
+
+def _fp8_upcast_linear_forward(self, input):
+    weight = self.weight.to(input.dtype) if self.weight.dtype in _FP8_DTYPES else self.weight
+    bias = self.bias.to(input.dtype) if (self.bias is not None and self.bias.dtype in _FP8_DTYPES) else self.bias
+    return torch.nn.functional.linear(input, weight, bias)
+
+
+def _patch_fp8_linears(model) -> int:
+    # This checkpoint ships its weights natively in fp8 (torch_dtype below preserves that
+    # instead of upcasting to bf16 at load time, halving resident memory: ~19GB vs ~38GB).
+    # nn.Linear has no fp8 compute kernel on this GPU, so each patched layer upcasts its own
+    # weight to the input's dtype just-in-time for the matmul — mathematically identical to
+    # the old load-time-upcast-everything approach (same values, same target dtype), just
+    # deferred so only one layer's weight is transiently bf16 at a time instead of all of them.
+    count = 0
+    for module in model.modules():
+        if isinstance(module, torch.nn.Linear) and module.weight.dtype in _FP8_DTYPES:
+            module.forward = types.MethodType(_fp8_upcast_linear_forward, module)
+            count += 1
+    return count
+
+
 _t0_load = time.perf_counter()
 print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
 _hb = _start_heartbeat("transformer")
 _transformer = QwenImageTransformer2DModel.from_pretrained(
     "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
-    torch_dtype=dtype,
+    torch_dtype=torch.float8_e4m3fn,
     device_map="cpu",
 )
 _hb.set()
 print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", flush=True)
+_n_fp8_patched = _patch_fp8_linears(_transformer)
+print(f"[startup] patched {_n_fp8_patched} fp8-resident nn.Linear modules for just-in-time upcast", flush=True)
+try:
+    print(f"[startup] transformer memory footprint: {_transformer.get_memory_footprint()/1024**3:.2f}GB", flush=True)
+except Exception as e:
+    print(f"[startup] transformer memory footprint: unavailable ({e})", flush=True)
 
 _t1_load = time.perf_counter()
 print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
@@ -404,13 +435,15 @@ def _log_gpu_properties(cuda_ok):
 
 # 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
+# transfer for several slower hook-managed ones. The previous bf16-everywhere
+# pipeline (~40GB transformer + ~14GB text encoder) peaked at 46.82GB moving
+# onto this Space's 47GB 2g.48gb MIG slice and still OOM'd, wasting ~40s
+# before falling back. The transformer now stays fp8-resident (see
+# _patch_fp8_linears above, ~19GB instead of ~38GB), so the full pipeline
+# should total roughly ~35GB — comfortably under the slice with headroom to
+# spare. Lowered accordingly, but the OOM fallback below stays as a safety
+# net in case that estimate is off.
+_FAST_PATH_MIN_GB = 40
 
 
 def _place_pipe_on_device(cuda_ok, gpu_props, t0):