Update app.py

app.py

@@ -81,20 +81,54 @@ from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
 
 from ltx_core.model.transformer import attention as _attn_mod
 
-print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
-try:
-    from xformers.ops import memory_efficient_attention as _mea
-    from xformers.ops.fmha import cutlass
-
-    def _cutlass_memory_efficient_attention(*args, **kwargs):
-        # Force CUTLASS and avoid FlashAttention paths that are crashing.
-        kwargs["op"] = (cutlass.FwOp, cutlass.BwOp)
-        return _mea(*args, **kwargs)
-
-    _attn_mod.memory_efficient_attention = _cutlass_memory_efficient_attention
-    print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
-except Exception as e:
-    print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")
+print(f"[ATTN] memory_efficient_attention={_attn_mod.memory_efficient_attention}")
+print(f"[ATTN] flash_attn_interface={_attn_mod.flash_attn_interface}")
+
+# If LTX-2's bare `import flash_attn_interface` failed, try the nested import
+# and inject it back so the built-in FlashAttention3 class works.
+if _attn_mod.flash_attn_interface is None:
+    try:
+        from flash_attn import flash_attn_interface as _flash_attn_interface
+        _attn_mod.flash_attn_interface = _flash_attn_interface
+        print("[ATTN] Recovered flash_attn_interface from flash_attn package")
+    except Exception as _e:
+        print(f"[ATTN] Could not recover flash_attn_interface: {type(_e).__name__}: {_e}")
+
+# ── Hard enforcement: error out if the wrong backend is present ──
+if _attn_mod.memory_efficient_attention is not None:
+    raise RuntimeError(
+        "xformers is still importable (memory_efficient_attention is not None). "
+        "Remove xformers from requirements.txt and rebuild the Space."
+    )
+
+if _attn_mod.flash_attn_interface is None:
+    raise RuntimeError(
+        "FlashAttention3 (flash_attn_interface) is not available. "
+        "Install flash-attn and ensure the Space has a Hopper-compatible build."
+    )
+# ──────────────────────────────────────────────────────────────────
+
+# Defensively wrap flash_attn_func: FA3 sometimes returns (output, softmax_lse)
+# while LTX-2 expects a single tensor.
+_fa3_func = _attn_mod.flash_attn_interface.flash_attn_func
+def _fa3_func_wrapped(q, k, v, *args, **kwargs):
+    result = _fa3_func(q, k, v, *args, **kwargs)
+    if isinstance(result, tuple):
+        return result[0]
+    return result
+
+_attn_mod.flash_attn_interface.flash_attn_func = _fa3_func_wrapped
+
+# Patch DEFAULT so it routes to FlashAttention3 instead of XFormers -> PyTorch
+_orig_attn_fn_call = _attn_mod.AttentionFunction.__call__
+
+def _default_to_fa3(self, q, k, v, heads, mask=None):
+    if self is _attn_mod.AttentionFunction.DEFAULT:
+        return _attn_mod.FlashAttention3()(q, k, v, heads, mask)
+    return _orig_attn_fn_call(self, q, k, v, heads, mask)
+
+_attn_mod.AttentionFunction.__call__ = _default_to_fa3
+print("[ATTN] Patched AttentionFunction.DEFAULT -> FlashAttention3")
 
 logging.getLogger().setLevel(logging.INFO)