add aoti for speed up

app.py

@@ -6,8 +6,7 @@ import sys
 os.environ["TORCH_COMPILE_DISABLE"] = "1"
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
 
-# memory-efficient attention
-subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
+# (removed runtime xformers install -> would pull torch 2.8 and break the AOTI .pt2; SDPA used)
 
 # --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
 LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
@@ -224,6 +223,38 @@ if not SKIP_STAGE_2:
     _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
 print("Pipeline ready (all components preloaded + pinned for ZeroGPU packing).")
 
+# ============================ AOTI (native bf16 transformer graph) ============================
+AOTI_REPO = os.environ.get("AOTI_REPO", "linoyts/LTX-2.3-Native-Transformer-GroupA-sm120-cu130-r20")
+import types as _types
+from dataclasses import replace as _dc_replace
+from ltx_core.model.transformer.transformer_args import TransformerArgs as _TA
+_TA_FIELDS = list(_TA.__dataclass_fields__.keys())
+def _flatten_ta(ta):
+    out = []
+    for f in _TA_FIELDS:
+        v = getattr(ta, f)
+        if torch.is_tensor(v):
+            out.append(v)
+        elif isinstance(v, tuple) and len(v) > 0 and all(torch.is_tensor(x) for x in v):
+            out.extend(v)
+    return out
+def _install_aoti():
+    velocity = pipeline.stage_1_model_ledger.transformer().velocity_model
+    spaces.aoti_load(module=velocity, repo_id=AOTI_REPO)
+    def _proc(self, video, audio, perturbations):
+        for blk in self.transformer_blocks:
+            o = blk(*(_flatten_ta(video) + _flatten_ta(audio)))
+            video = _dc_replace(video, x=o[0]); audio = _dc_replace(audio, x=o[1])
+        return video, audio
+    velocity._process_transformer_blocks = _types.MethodType(_proc, velocity)
+    print(f"[AOTI] loaded {AOTI_REPO} + patched block loop", flush=True)
+print(f"[AOTI] base torch={torch.__version__} cuda={torch.version.cuda}", flush=True)
+try:
+    _install_aoti(); print("[AOTI] OK", flush=True)
+except Exception as _e:
+    import traceback; traceback.print_exc(); print(f"[AOTI] FAILED ({_e!r}) -> EAGER", flush=True)
+# ==============================================================================================
+
 
 def _duration(*args, **kwargs):
     nf = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), DEFAULT_FRAMES)
@@ -270,6 +301,7 @@ with gr.Blocks(title="LTX-2.3 Colorize") as demo:
         "motion identity. Using [LTX 2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) "
         "with the [Colorization IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Colorization)."
     )
+    gr.Markdown("⚡ **Accelerated with [AOTI](https://huggingface.co/linoyts/LTX-2.3-Native-Transformer-GroupA-sm120-cu130-r20)** — precompiled transformer for faster inference.")
     with gr.Row():
         with gr.Column():
             video_in = gr.Video(label="Input video (any clip — recolored as B&W)")

requirements.txt

@@ -1,11 +1,12 @@
+--extra-index-url https://download.pytorch.org/whl/cu130
 transformers==4.57.6
 accelerate
-torch==2.8.0
-torchaudio==2.8.0
+torchaudio==2.11.0
 einops
 scipy
 av
-scikit-image>=0.25.2
-flashpack==0.1.2
 imageio[ffmpeg]
 pillow
+spaces
+sentencepiece
+ftfy