Update app.py

app.py

@@ -2,10 +2,6 @@ import os
 import subprocess
 import sys
 
-# Disable torch.compile / dynamo before any torch import
-os.environ["TORCH_COMPILE_DISABLE"] = "1"
-os.environ["TORCHDYNAMO_DISABLE"] = "1"
-
 # 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")
@@ -36,8 +32,8 @@ import gc
 import hashlib
 
 import torch
-torch._dynamo.config.suppress_errors = True
-torch._dynamo.config.disable = True
+torch._dynamo.config.suppress_errors = False
+torch._dynamo.config.disable = False
 
 import spaces
 import gradio as gr
@@ -700,6 +696,18 @@ print("=" * 80)
 print("Pipeline ready!")
 print("=" * 80)
 
+# AFTER your preload block, compile the transformer
+print("Compiling transformer with torch.compile...")
+
+# Regional compilation on transformer_blocks is best for DiT models
+# This compiles the repeated BasicAVTransformerBlock layers and reuses the graph
+_transformer = torch.compile(
+    _transformer,
+    mode="max-autotune",
+    fullgraph=False,  # safer; set True if no graph breaks
+    dynamic=False,    # critical: must be False for fixed shapes per generation
+)
+
 def log_memory(tag: str):
     if torch.cuda.is_available():
         allocated = torch.cuda.memory_allocated() / 1024**3