Upload app.py with huggingface_hub

app.py

@@ -282,16 +282,18 @@ def _edit_transformer_forward(latents, src_packed, prompt_embeds, prompt_mask,
 
     image_rotary_emb = m.rotary_emb(position_ids)
 
-    # Attention mask: text uses its key-padding mask, all image tokens (src+tgt) valid.
+    # Single prompt, batch=1: the grounded encode uses padding=True on one
+    # sequence, so there is no padding and every token is valid. Passing
+    # attention_mask=None lets SDPA pick the flash / mem-efficient kernels
+    # (a dense boolean mask would force the score-materializing math path and
+    # OOM given the doubled token count). If padding is ever introduced (batched
+    # or CFG with unequal lengths), trim to the shorter length instead of masking.
     attention_mask = None
-    if prompt_mask is not None:
-        img_ones = prompt_mask.new_ones((combined_img.shape[0], combined_img.shape[1]))
-        attention_mask = torch.cat([prompt_mask, img_ones], dim=1)[:, None, None, :]
-
-    # Force the math SDPA kernel. The flash / mem-efficient kernels probe the
-    # device via NVML, which on ZeroGPU's forked CUDA context triggers
-    # "NVML_SUCCESS == r INTERNAL ASSERT FAILED" inside the caching allocator.
-    with sdpa_kernel([SDPBackend.MATH]):
+
+    # Prefer the flash / mem-efficient SDPA kernels (they never materialize the
+    # full seq x seq score matrix, which the MATH kernel would OOM on given the
+    # doubled token count from prepending the source latent).
+    with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
         for block in m.transformer_blocks:
             hidden = block(hidden, temb_mod, image_rotary_emb, attention_mask)