Update app.py

app.py

@@ -6,11 +6,6 @@ import sys
 os.environ["TORCH_COMPILE_DISABLE"] = "1"
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
 
-subprocess.run(
-    [sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
-    check=False,
-)
-
 # Clone LTX-2 repo and install packages
 LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
 LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
@@ -83,34 +78,22 @@ from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
 from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
 from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
 
-from flash_attn import flash_attn_func
+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 ltx_core.model.transformer import attention as _attn_mod
-    _attn_mod.memory_efficient_attention = flash_attn_func
-except Exception as e:
-    print(f"[ATTN] FA3 import failed: {e}, using xformers")
     from xformers.ops import memory_efficient_attention as _mea
-    _attn_mod.memory_efficient_attention = _mea
+    from xformers.ops.fmha import cutlass
 
-# Attention kernels: try FlashAttention-3 first, fall back to xformers
-from ltx_core.model.transformer import attention as _attn_mod
-print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
+    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)
 
-try:
-    from kernels import get_kernel
-    fa3_kernel = get_kernel("kernels-community/vllm-flash-attn3")
-    _attn_mod.memory_efficient_attention = fa3_kernel
-    print(f"[ATTN] FA3 kernel applied: {_attn_mod.memory_efficient_attention}")
+    _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] FA3 not available ({type(e).__name__}: {e}), trying xformers...")
-
-    try:
-        from xformers.ops import memory_efficient_attention as _mea
-        _attn_mod.memory_efficient_attention = _mea
-        print(f"[ATTN] xformers applied: {_attn_mod.memory_efficient_attention}")
-    except Exception as e2:
-        print(f"[ATTN] xformers also FAILED: {type(e2).__name__}: {e2}")
+    print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")
 
 logging.getLogger().setLevel(logging.INFO)