Update app.py

app.py

@@ -78,22 +78,24 @@ 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
 
+# 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}")
-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}")
+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}")
 except Exception as e:
-    print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {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}")
 
 logging.getLogger().setLevel(logging.INFO)
 
@@ -289,8 +291,7 @@ class LTX23DistilledA2VPipeline:
         )
 
         torch.cuda.synchronize()
-        del video_state       # release Stage 1 latent only, keep encoder/transformer on GPU
-        torch.cuda.empty_cache()  # lightweight, no sync
+        cleanup_memory()
 
         # ── Upscaling ──
         upscaled_video_latent = upsample_video(
@@ -326,6 +327,7 @@ class LTX23DistilledA2VPipeline:
         )
 
         torch.cuda.synchronize()
+        cleanup_memory()
 
         # ── Decode both video and audio ──
         decoded_video = vae_decode_video(