SDPA fallback for the CLIP/DiT flash_attention path; Gradio 6 theme in launch()

app.py

@@ -407,7 +407,7 @@ CSS = """
 .gradio-container { max-width: 1200px !important; }
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="ForgeWM") as demo:
+with gr.Blocks(title="ForgeWM") as demo:
     gr.Markdown(
         "# 🎮 ForgeWM — few-step action-conditioned world model\n"
         "Give it **one Minecraft frame** and a short **action script**; the "
@@ -480,4 +480,5 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="ForgeWM") as demo:
         api_name="generate",
     )
 
-demo.queue(max_size=12).launch(mcp_server=True, show_error=True)
+demo.queue(max_size=12).launch(
+    theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)

wan/modules/attention.py

@@ -26,6 +26,7 @@ import warnings
 __all__ = [
     'flash_attention',
     'attention',
+    'sdpa_flash_attn_func',
 ]
 
 
@@ -57,6 +58,23 @@ def flash_attention(
     deterministic:  bool. If True, slightly slower and uses more memory.
     dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
     """
+    if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE):
+        # No FlashAttention wheel on this platform.  Fall back to PyTorch SDPA
+        # while preserving the caller's dtype (the FA path would downcast to
+        # bfloat16, which we do not want for the fp32 CLIP tower).
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                'Padding mask is disabled when using scaled_dot_product_attention.')
+        if q_scale is not None:
+            q = q * q_scale
+        return sdpa_flash_attn_func(
+            q, k, v,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+        )
+
     half_dtypes = (torch.float16, torch.bfloat16)
     assert dtype in half_dtypes
     assert q.device.type == 'cuda' and q.size(-1) <= 256