linoyts HF Staff commited on
Commit
cd7fe6c
·
1 Parent(s): 489e53f

add aoti for speed up (#5)

Browse files

- AOTI: faster transformer (native bf16, ZeroGPU) (ef888cce6dc11b2ef398bcd2542a28159a4bd165)
- add aoti for speed up + mention in description (e1ce7b4c8a07ff5a7a7ddf62cda448abea9f748a)

Files changed (2) hide show
  1. app.py +34 -2
  2. requirements.txt +5 -4
app.py CHANGED
@@ -6,8 +6,7 @@ import sys
6
  os.environ["TORCH_COMPILE_DISABLE"] = "1"
7
  os.environ["TORCHDYNAMO_DISABLE"] = "1"
8
 
9
- # memory-efficient attention
10
- subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
11
 
12
  # --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
13
  LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
@@ -210,6 +209,38 @@ if not SKIP_STAGE_2:
210
  _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
211
  print("Pipeline ready.")
212
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
213
 
214
  def _sheet_to_video(img, width, height, num_frames):
215
  """Repeat the reference-sheet image into an NF-frame video for video_conditioning."""
@@ -301,6 +332,7 @@ with gr.Blocks(title="LTX-2.3 Ingredients (Fast)") as demo:
301
  "[LTX 2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) with the "
302
  "[Ingredients IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Ingredients). "
303
  )
 
304
  with gr.Row():
305
  with gr.Column():
306
  with gr.Tabs():
 
6
  os.environ["TORCH_COMPILE_DISABLE"] = "1"
7
  os.environ["TORCHDYNAMO_DISABLE"] = "1"
8
 
9
+ # (removed runtime xformers install -> would pull torch 2.8 and break the AOTI .pt2; SDPA used)
 
10
 
11
  # --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
12
  LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
 
209
  _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
210
  print("Pipeline ready.")
211
 
212
+ # ============================ AOTI (native bf16 transformer graph) ============================
213
+ AOTI_REPO = os.environ.get("AOTI_REPO", "linoyts/LTX-2.3-Native-Transformer-GroupA-sm120-cu130-r20")
214
+ import types as _types
215
+ from dataclasses import replace as _dc_replace
216
+ from ltx_core.model.transformer.transformer_args import TransformerArgs as _TA
217
+ _TA_FIELDS = list(_TA.__dataclass_fields__.keys())
218
+ def _flatten_ta(ta):
219
+ out = []
220
+ for f in _TA_FIELDS:
221
+ v = getattr(ta, f)
222
+ if torch.is_tensor(v):
223
+ out.append(v)
224
+ elif isinstance(v, tuple) and len(v) > 0 and all(torch.is_tensor(x) for x in v):
225
+ out.extend(v)
226
+ return out
227
+ def _install_aoti():
228
+ velocity = pipeline.stage_1_model_ledger.transformer().velocity_model
229
+ spaces.aoti_load(module=velocity, repo_id=AOTI_REPO)
230
+ def _proc(self, video, audio, perturbations):
231
+ for blk in self.transformer_blocks:
232
+ o = blk(*(_flatten_ta(video) + _flatten_ta(audio)))
233
+ video = _dc_replace(video, x=o[0]); audio = _dc_replace(audio, x=o[1])
234
+ return video, audio
235
+ velocity._process_transformer_blocks = _types.MethodType(_proc, velocity)
236
+ print(f"[AOTI] loaded {AOTI_REPO} + patched block loop", flush=True)
237
+ print(f"[AOTI] base torch={torch.__version__} cuda={torch.version.cuda}", flush=True)
238
+ try:
239
+ _install_aoti(); print("[AOTI] OK", flush=True)
240
+ except Exception as _e:
241
+ import traceback; traceback.print_exc(); print(f"[AOTI] FAILED ({_e!r}) -> EAGER", flush=True)
242
+ # ==============================================================================================
243
+
244
 
245
  def _sheet_to_video(img, width, height, num_frames):
246
  """Repeat the reference-sheet image into an NF-frame video for video_conditioning."""
 
332
  "[LTX 2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) with the "
333
  "[Ingredients IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Ingredients). "
334
  )
335
+ gr.Markdown("⚡ **Accelerated with [AOTI](https://huggingface.co/linoyts/LTX-2.3-Native-Transformer-GroupA-sm120-cu130-r20)** — precompiled transformer for faster inference.")
336
  with gr.Row():
337
  with gr.Column():
338
  with gr.Tabs():
requirements.txt CHANGED
@@ -1,11 +1,12 @@
 
1
  transformers==4.57.6
2
  accelerate
3
- torch==2.8.0
4
- torchaudio==2.8.0
5
  einops
6
  scipy
7
  av
8
- scikit-image>=0.25.2
9
- flashpack==0.1.2
10
  imageio[ffmpeg]
11
  pillow
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu130
2
  transformers==4.57.6
3
  accelerate
4
+ torchaudio==2.11.0
 
5
  einops
6
  scipy
7
  av
 
 
8
  imageio[ffmpeg]
9
  pillow
10
+ spaces
11
+ sentencepiece
12
+ ftfy