• Stefy Lanza (nextime / spora )'s avatar
    fix: Z-Image image-gen 500 — flash-attn-2 rejects attn_mask (shared-backend race) · 643fa155
    Stefy Lanza (nextime / spora ) authored
    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: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
    Claude-Session: https://claude.ai/code/session_01LoSpEthysqmseCc6Geizty
    643fa155
images.py 110 KB