Load text_encoder as int8 (bitsandbytes) when TEXT_ENCODER_INT8_REPO is set

bitsandbytes 8bit models can't be moved between devices with `.to()`
(unconditional in transformers, unlike the version-gated allowance for
4bit), so this can't reuse the transformer's cpu-load-then-.to(device)
pattern. Instead the quantized text_encoder is loaded directly onto the
target CUDA device on first inference call per worker, guarded so it
only happens once. Also fixes a crash: high_detail mode's pre-decode
text_encoder.to("cpu") offload is skipped once text_encoder is int8,
since .to() would raise for it regardless of target device.

Repo id comes from the TEXT_ENCODER_INT8_REPO secret; falls back to the
existing bf16 text_encoder when unset.

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

app.py

@@ -75,6 +75,7 @@ from mode import Mode
 print("[startup] importing diffusers...", flush=True)
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.models.normalization import RMSNorm
+from transformers import Qwen2_5_VLForConditionalGeneration
 print("[startup] importing QwenImageEditPlusPipeline...", flush=True)
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 print("[startup] importing QwenImageTransformer2DModel...", flush=True)
@@ -501,6 +502,37 @@ def _log_gpu_properties(cuda_ok):
 # net in case that estimate is off.
 _FAST_PATH_MIN_GB = 40
 
+# int8 (bitsandbytes) quantized text_encoder, produced offline from the FireRed
+# text_encoder's bf16 weights (~8.75GB vs ~15.4GB for the bf16 original). Repo id comes
+# from a secret rather than being hardcoded here.
+_TEXT_ENCODER_INT8_REPO = os.environ.get("TEXT_ENCODER_INT8_REPO")
+
+
+def _ensure_int8_text_encoder(cuda_ok, t0):
+    # 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
+    # can't follow the fp8 transformer's cpu-load-then-.to(device) pattern — it must be loaded
+    # directly onto the target CUDA device, which is only visible inside this @spaces.GPU call.
+    # Guarded so it only runs once per worker; diffusers' pipe.to(device) in
+    # _place_pipe_on_device already knows to skip an 8bit-quantized module it finds pre-placed.
+    if not cuda_ok or not _TEXT_ENCODER_INT8_REPO or getattr(pipe, "_text_encoder_is_int8", False):
+        return
+    try:
+        _t_load = time.perf_counter()
+        quantized = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            _TEXT_ENCODER_INT8_REPO,
+            device_map={"": device},
+            dtype=torch.bfloat16,
+        )
+        pipe.text_encoder = quantized
+        pipe._text_encoder_is_int8 = True
+        print(
+            f"[infer] loaded int8 text_encoder from {_TEXT_ENCODER_INT8_REPO} — "
+            f"{(time.perf_counter()-_t_load)*1000:.0f}ms | t={time.perf_counter()-t0:.1f}s"
+        )
+    except Exception as e:
+        print(f"[infer] WARNING: int8 text_encoder load failed, keeping bf16: {type(e).__name__}: {e}")
+
 
 def _place_pipe_on_device(cuda_ok, gpu_props, t0):
     if getattr(pipe.transformer, "_hf_hook", None) is not None:
@@ -546,14 +578,18 @@ def _make_step_callback(steps, timer, t0, mode: Mode):
         # 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.
+        # to avoid fighting their own device bookkeeping, and when text_encoder is int8
+        # (bitsandbytes) quantized — `.to()` is unconditionally unsupported for 8bit models
+        # (would raise), and its ~8.75GB footprint needs this safety net less anyway.
         if step_idx == steps - 1 and getattr(pipeline.text_encoder, "_hf_hook", None) is None:
-            if mode.offloads_text_encoder_before_decode:
+            if mode.offloads_text_encoder_before_decode and not getattr(pipeline, "_text_encoder_is_int8", False):
                 _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")
+            elif getattr(pipeline, "_text_encoder_is_int8", False):
+                print("[infer] skipping text_encoder offload (int8, .to() unsupported / smaller footprint)")
             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
@@ -582,6 +618,7 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
     _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode)
     gpu_props = _log_gpu_properties(_cuda_ok)
 
+    _ensure_int8_text_encoder(_cuda_ok, t0)
     _place_pipe_on_device(_cuda_ok, gpu_props, t0)
     print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
 

requirements.txt

@@ -2,6 +2,7 @@ accelerate==1.14.0
 diffusers==0.39.0
 peft==0.19.1
 transformers==4.57.1
+bitsandbytes==0.47.0
 huggingface_hub
 pyarrow
 sentencepiece