Quantize transformer to NF4 via bitsandbytes to fix root-cause OOM

The last several commits kept rearranging when the ~54GB bf16 pipeline
(20B transformer + Qwen2.5-VL-7B text encoder) touches the GPU, but that
footprint never comfortably fit the 47GB 2g.48gb MIG slice regardless of
ordering. NF4-quantizing the transformer (the biggest contributor) drops
it from ~40GB to ~11GB, bringing full-pipeline GPU residency down to
~25-30GB with real headroom, so the fast pipe.to(device) path should now
succeed every time instead of needing the OOM/cpu-offload dance.

bitsandbytes rejects device_map="cpu" for a fresh (non-prequantized)
config, so the transformer now loads straight onto cuda:0 at startup
(paying the "ZeroGPU disk packing" cost once at boot) instead of the
CPU-first pattern still used for the rest of the pipeline. Removed the
now-obsolete _FAST_PATH_MIN_GB size gate in _infer_gpu since the
quantized pipeline should fit regardless of the exact slice size.

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

app.py

@@ -72,7 +72,7 @@ print("[startup] TF32 enabled", flush=True)
 print("[startup] importing dimensions...", flush=True)
 from dimensions import compute_output_dimensions, max_dim_for_mode
 print("[startup] importing diffusers...", flush=True)
-from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers import BitsAndBytesConfig, FlowMatchEulerDiscreteScheduler
 print("[startup] importing QwenImageEditPlusPipeline...", flush=True)
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 print("[startup] importing QwenImageTransformer2DModel...", flush=True)
@@ -97,13 +97,31 @@ def _start_heartbeat(label: str) -> threading.Event:
 _t0_load = time.perf_counter()
 print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
 _hb = _start_heartbeat("transformer")
+# NF4-quantize the transformer (~20B params, ~40GB in bf16) to fit the pipeline
+# comfortably inside this Space's 47GB MIG slice. bitsandbytes only quantizes
+# on an actual CUDA device (it rejects device_map="cpu"/"disk" for a fresh,
+# not-yet-quantized config), so this loads straight onto cuda:0 instead of the
+# CPU-first pattern used below for the rest of the pipeline. That pays the
+# once-per-boot "ZeroGPU disk packing" cost the CPU-first load was added to
+# avoid, but only once at startup rather than never.
+_quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=dtype,
+    bnb_4bit_use_double_quant=True,
+) if torch.cuda.is_available() else None
 _transformer = QwenImageTransformer2DModel.from_pretrained(
     "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
     torch_dtype=dtype,
-    device_map="cpu",
+    device_map={"": 0} if _quant_config else "cpu",
+    quantization_config=_quant_config,
 )
 _hb.set()
 print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", 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,27 +422,21 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
 
     # 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
+    # to(device) transfer for several slower hook-managed ones. The
+    # transformer is now NF4-quantized (~40GB bf16 -> ~11GB), so the full
+    # pipeline (~25-30GB total) should comfortably fit even on the 47GB
+    # 2g.48gb MIG slice this Space actually gets — always attempt full-GPU
+    # residency first and only fall back to offload on an actual OOM.
     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:
+        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} (slice too small for fast path)")
+            print(f"[infer] enabled cpu offload on {device} (fallback)")
     print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
 
     print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")

requirements.txt

@@ -2,6 +2,7 @@ git+https://github.com/huggingface/accelerate.git
 git+https://github.com/huggingface/diffusers.git
 git+https://github.com/huggingface/peft.git
 transformers==4.57.1
+bitsandbytes
 huggingface_hub
 pyarrow
 sentencepiece