dagloop5 commited on
Commit
2430250
·
verified ·
1 Parent(s): 07cf967

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -6,9 +6,6 @@ import sys
6
  os.environ["TORCH_COMPILE_DISABLE"] = "1"
7
  os.environ["TORCHDYNAMO_DISABLE"] = "1"
8
 
9
- # Install xformers for memory-efficient attention
10
- subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
11
-
12
  # Clone LTX-2 repo and install packages
13
  LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
14
  LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
@@ -90,10 +87,18 @@ from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
90
 
91
  # Force-patch xformers attention into the LTX attention module.
92
  from ltx_core.model.transformer import attention as _attn_mod
 
93
  print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
94
  try:
95
  from xformers.ops import memory_efficient_attention as _mea
96
- _attn_mod.memory_efficient_attention = _mea
 
 
 
 
 
 
 
97
  print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
98
  except Exception as e:
99
  print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")
 
6
  os.environ["TORCH_COMPILE_DISABLE"] = "1"
7
  os.environ["TORCHDYNAMO_DISABLE"] = "1"
8
 
 
 
 
9
  # Clone LTX-2 repo and install packages
10
  LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
11
  LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
 
87
 
88
  # Force-patch xformers attention into the LTX attention module.
89
  from ltx_core.model.transformer import attention as _attn_mod
90
+
91
  print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
92
  try:
93
  from xformers.ops import memory_efficient_attention as _mea
94
+ from xformers.ops.fmha import cutlass
95
+
96
+ def _cutlass_memory_efficient_attention(*args, **kwargs):
97
+ # Force CUTLASS and avoid FlashAttention paths that are crashing.
98
+ kwargs["op"] = (cutlass.FwOp, cutlass.BwOp)
99
+ return _mea(*args, **kwargs)
100
+
101
+ _attn_mod.memory_efficient_attention = _cutlass_memory_efficient_attention
102
  print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
103
  except Exception as e:
104
  print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")