fix: Z-Image image-gen 500 — flash-attn-2 rejects attn_mask (shared-backend race)

POST /v1/images/generations for a masked image transformer (Z-Image) crashed
with "`attn_mask` is not supported for flash-attn 2." → HTTP 500.

The diffusers "active attention backend" is PROCESS-WIDE global mutable state
(_AttentionBackendRegistry class attribute) shared by the image and video paths.
The video path sets it to flash-attn-2. A transformer whose per-module backend
is None reads that global at dispatch time; flash-attn-2 rejects attn_mask, so
Z-Image (which uses caption masks) crashes. The previous fix reset the global to
native before generating, but that reset is NOT atomic with the pipeline() call
(it runs in a to_thread worker), so a concurrent/subsequent video generation
races and flips the shared global back to flash before image attention dispatches.

Fix: pin the image denoiser's per-module attention backend to native (SDPA,
mask-supporting) via set_attention_backend("native"). The dispatcher passes
processor._attention_backend explicitly, bypassing the shared global — so image
attention is immune to the video path flipping it. Video is undisturbed (it sets
its own explicit per-module 'flash'). Image pipelines were never intentionally on
flash, so this is non-regressive. Global reset kept as a fallback.

Verified: unsloth/Z-Image-Turbo-unsloth-bnb-4bit now returns 200 with an image;
no "attn_mask is not supported" in the log.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LoSpEthysqmseCc6Geizty

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.1.62"
+__version__ = "0.1.63"
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
 # expandable_segments lets the allocator return freed pages to the driver even

codai/api/images.py

@@ -975,14 +975,30 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
         except Exception as _ip_err:
             print(f"Warning: IP-Adapter injection failed ({_ip_err}), continuing without character refs")
 
-    # Reset the diffusers GLOBAL attention backend to the environment default
-    # (native/SDPA) before generating. diffusers' Model.set_attention_backend()
-    # ALSO flips a process-wide active backend, and the video path sets it to
-    # flash-attn — which then leaks to image transformers that don't set their own
-    # (e.g. Z-Image passes backend=None → uses the global) and crashes with
-    # "`attn_mask` is not supported for flash-attn 2". Image + video share the
-    # engine process, so restore the default here so masked image attention (SDPA)
-    # always works. Cheap + idempotent; no-op if diffusers lacks the dispatcher.
+    # Pin the denoising model's attention to a mask-supporting backend before
+    # generating. The diffusers "active attention backend" is PROCESS-WIDE global
+    # mutable state (a class attribute on _AttentionBackendRegistry) shared with the
+    # video path, which sets it to flash-attn-2. A transformer whose per-module
+    # backend is None reads that global at dispatch time — and flash-attn-2 REJECTS
+    # attn_mask ("`attn_mask` is not supported for flash-attn 2."), so a masked image
+    # transformer (e.g. Z-Image) crashes → HTTP 500. Simply resetting the global
+    # isn't enough: the reset isn't atomic with the pipeline() call (which runs in a
+    # to_thread worker), so a concurrent/subsequent video generation can flip the
+    # shared global back to flash before image attention dispatches — a race.
+    #
+    # Setting the per-module backend EXPLICITLY makes image dispatch immune to the
+    # shared global (the dispatcher passes processor._attention_backend, bypassing
+    # get_active_backend()). native (SDPA) supports masks. This doesn't disturb the
+    # video path — it sets its own explicit per-module 'flash', so it ignores the
+    # global too. Image pipelines were never intentionally on flash, so pinning
+    # native here is non-regressive. Cheap + idempotent; no-op if unsupported.
+    for _cn in ("transformer", "unet"):
+        _c = getattr(pipeline, _cn, None)
+        if _c is not None and hasattr(_c, "set_attention_backend"):
+            try:
+                _c.set_attention_backend("native")
+            except Exception:
+                pass
     try:
         from diffusers.models.attention_dispatch import (
             _AttentionBackendRegistry, AttentionBackendName)