Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+package/submodules/transformer_blocks/package.pt2 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,161 @@
+
+---
+tags:
+- ahead-of-time
+- pytorch
+library_name: diffusers
+---
+
+> [!NOTE]
+> This **README** has been auto-generated by the **HF Job** run linked below
+> and the whole repository is a reproducible artifact of this Job
+
+# Ahead-of-time repository
+
+AoT repos contain **pre-compiled binaries** of PyTorch models, enabling:
+- fast startup times (no `torch.compile` needed)
+- significant **speedup**
+- **ZeroGPU** compatibility
+
+## How to use
+``` python
+
+import os
+import tempfile
+
+import numpy as np
+import torch
+import spaces
+from PIL import Image
+
+from diffusers import LTX2InContextPipeline
+from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
+from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
+
+# base == distilled in architecture, so one compiled graph serves both; distilled is
+# what most demos use. The AOTI package is weight-agnostic, so this base graph also
+# serves any FUSED LoRA (fuse_lora before aoti_load on the Space).
+MODEL_ID = os.environ.get("LTX_MODEL_ID", "diffusers/LTX-2.3-Distilled-Diffusers")
+
+pipe = LTX2InContextPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe.vae.enable_tiling()
+
+spaces.aoti_load(
+    module=pipe.transformer,
+    repo_id='linoyts/LTX-2.3-Transformer-GroupC-STG-sm120-cu130-rb3',
+)
+```
+
+## How to reproduce or customize
+``` bash
+hf jobs uv run job.py --flavor rtx-pro-6000 --image pytorch/pytorch:2.9.1-cuda13.0-cudnn9-devel --secrets HF_TOKEN
+```
+
+## Samples
+| Before compilation (0.42s) | After compilation (0.33s) |
+|---|---|
+| <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupC-STG-sm120-cu130-rb3/resolve/main/samples/before/video.mp4" controls></video> | <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupC-STG-sm120-cu130-rb3/resolve/main/samples/after/video.mp4" controls></video> |
+
+Speedup: **1.27x**
+
+## Environment
+<details><summary>Click to expand</summary>
+
+```
+PyTorch version: 2.12.0+cu130
+Is debug build: False
+CUDA used to build PyTorch: 13.0
+ROCM used to build PyTorch: N/A
+
+OS: Ubuntu 22.04.5 LTS (x86_64)
+GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
+Clang version: Could not collect
+CMake version: version 4.1.2
+Libc version: glibc-2.35
+
+Python version: 3.10.19 (main, Oct 31 2025, 23:02:46) [Clang 21.1.4 ] (64-bit runtime)
+Python platform: Linux-6.12.88-119.157.amzn2023.x86_64-x86_64-with-glibc2.35
+Is CUDA available: True
+CUDA runtime version: 13.0.48
+CUDA_MODULE_LOADING set to: 
+GPU models and configuration: GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition
+Nvidia driver version: 580.159.03
+cuDNN version: Could not collect
+Is XPU available: False
+HIP runtime version: N/A
+MIOpen runtime version: N/A
+Is XNNPACK available: True
+Caching allocator config: N/A
+
+CPU:
+Architecture:                            x86_64
+CPU op-mode(s):                          32-bit, 64-bit
+Address sizes:                           46 bits physical, 48 bits virtual
+Byte Order:                              Little Endian
+CPU(s):                                  192
+On-line CPU(s) list:                     0-191
+Vendor ID:                               GenuineIntel
+Model name:                              Intel(R) Xeon(R) Platinum 8559C
+CPU family:                              6
+Model:                                   207
+Thread(s) per core:                      2
+Core(s) per socket:                      48
+Socket(s):                               2
+Stepping:                                2
+BogoMIPS:                                4800.00
+Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd ida arat avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid cldemote movdiri movdir64b md_clear serialize amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities
+Hypervisor vendor:                       KVM
+Virtualization type:                     full
+L1d cache:                               4.5 MiB (96 instances)
+L1i cache:                               3 MiB (96 instances)
+L2 cache:                                192 MiB (96 instances)
+L3 cache:                                640 MiB (2 instances)
+NUMA node(s):                            2
+NUMA node0 CPU(s):                       0-47,96-143
+NUMA node1 CPU(s):                       48-95,144-191
+Vulnerability Gather data sampling:      Not affected
+Vulnerability Indirect target selection: Not affected
+Vulnerability Itlb multihit:             Not affected
+Vulnerability L1tf:                      Not affected
+Vulnerability Mds:                       Not affected
+Vulnerability Meltdown:                  Not affected
+Vulnerability Mmio stale data:           Not affected
+Vulnerability Reg file data sampling:    Not affected
+Vulnerability Retbleed:                  Not affected
+Vulnerability Spec rstack overflow:      Not affected
+Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
+Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
+Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
+Vulnerability Srbds:                     Not affected
+Vulnerability Tsa:                       Not affected
+Vulnerability Tsx async abort:           Not affected
+Vulnerability Vmscape:                   Not affected
+
+Versions of relevant libraries:
+[pip3] Could not collect
+[conda] numpy                        2.3.4            py311h2e04523_0       conda-forge
+[conda] nvidia-cublas                13.0.0.19        pypi_0                pypi
+[conda] nvidia-cuda-cupti            13.0.48          pypi_0                pypi
+[conda] nvidia-cuda-nvrtc            13.0.48          pypi_0                pypi
+[conda] nvidia-cuda-runtime          13.0.48          pypi_0                pypi
+[conda] nvidia-cudnn-cu13            9.13.0.50        pypi_0                pypi
+[conda] nvidia-cufft                 12.0.0.15        pypi_0                pypi
+[conda] nvidia-curand                10.4.0.35        pypi_0                pypi
+[conda] nvidia-cusolver              12.0.3.29        pypi_0                pypi
+[conda] nvidia-cusparse              12.6.2.49        pypi_0                pypi
+[conda] nvidia-cusparselt-cu13       0.8.0            pypi_0                pypi
+[conda] nvidia-nccl-cu13             2.27.7           pypi_0                pypi
+[conda] nvidia-nvjitlink             13.0.39          pypi_0                pypi
+[conda] nvidia-nvtx                  13.0.39          pypi_0                pypi
+[conda] optree                       0.17.0           pypi_0                pypi
+[conda] torch                        2.9.1+cu130      pypi_0                pypi
+[conda] torchaudio                   2.9.1+cu130      pypi_0                pypi
+[conda] torchelastic                 0.2.2            pypi_0                pypi
+[conda] torchvision                  0.24.1+cu130     pypi_0                pypi
+[conda] triton                       3.5.1            pypi_0                pypi
+```
+</details>
+
+## Job run
+- [linoyts/6a3288ccfb114ff24a3885f5](https://huggingface.co/jobs/linoyts/6a3288ccfb114ff24a3885f5)

environment.json

@@ -0,0 +1,28 @@
+{
+    "torch_version": "2.12.0+cu130",
+    "is_debug_build": "False",
+    "cuda_compiled_version": "13.0",
+    "gcc_version": "(Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0",
+    "clang_version": null,
+    "cmake_version": "version 4.1.2",
+    "os": "Ubuntu 22.04.5 LTS (x86_64)",
+    "libc_version": "glibc-2.35",
+    "python_version": "3.10.19 (main, Oct 31 2025, 23:02:46) [Clang 21.1.4 ] (64-bit runtime)",
+    "python_platform": "Linux-6.12.88-119.157.amzn2023.x86_64-x86_64-with-glibc2.35",
+    "is_cuda_available": "True",
+    "cuda_runtime_version": "13.0.48",
+    "cuda_module_loading": "",
+    "nvidia_driver_version": "580.159.03",
+    "nvidia_gpu_models": "GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition",
+    "cudnn_version": null,
+    "is_xpu_available": "False",
+    "pip_version": "pip3",
+    "pip_packages": null,
+    "conda_packages": "numpy                        2.3.4            py311h2e04523_0       conda-forge\nnvidia-cublas                13.0.0.19        pypi_0                pypi\nnvidia-cuda-cupti            13.0.48          pypi_0                pypi\nnvidia-cuda-nvrtc            13.0.48          pypi_0                pypi\nnvidia-cuda-runtime          13.0.48          pypi_0                pypi\nnvidia-cudnn-cu13            9.13.0.50        pypi_0                pypi\nnvidia-cufft                 12.0.0.15        pypi_0                pypi\nnvidia-curand                10.4.0.35        pypi_0                pypi\nnvidia-cusolver              12.0.3.29        pypi_0                pypi\nnvidia-cusparse              12.6.2.49        pypi_0                pypi\nnvidia-cusparselt-cu13       0.8.0            pypi_0                pypi\nnvidia-nccl-cu13             2.27.7           pypi_0                pypi\nnvidia-nvjitlink             13.0.39          pypi_0                pypi\nnvidia-nvtx                  13.0.39          pypi_0                pypi\noptree                       0.17.0           pypi_0                pypi\ntorch                        2.9.1+cu130      pypi_0                pypi\ntorchaudio                   2.9.1+cu130      pypi_0                pypi\ntorchelastic                 0.2.2            pypi_0                pypi\ntorchvision                  0.24.1+cu130     pypi_0                pypi\ntriton                       3.5.1            pypi_0                pypi",
+    "hip_compiled_version": "N/A",
+    "hip_runtime_version": "N/A",
+    "miopen_runtime_version": "N/A",
+    "caching_allocator_config": {},
+    "is_xnnpack_available": "True",
+    "cpu_info": "Architecture:                            x86_64\nCPU op-mode(s):                          32-bit, 64-bit\nAddress sizes:                           46 bits physical, 48 bits virtual\nByte Order:                              Little Endian\nCPU(s):                                  192\nOn-line CPU(s) list:                     0-191\nVendor ID:                               GenuineIntel\nModel name:                              Intel(R) Xeon(R) Platinum 8559C\nCPU family:                              6\nModel:                                   207\nThread(s) per core:                      2\nCore(s) per socket:                      48\nSocket(s):                               2\nStepping:                                2\nBogoMIPS:                                4800.00\nFlags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd ida arat avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid cldemote movdiri movdir64b md_clear serialize amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities\nHypervisor vendor:                       KVM\nVirtualization type:                     full\nL1d cache:                               4.5 MiB (96 instances)\nL1i cache:                               3 MiB (96 instances)\nL2 cache:                                192 MiB (96 instances)\nL3 cache:                                640 MiB (2 instances)\nNUMA node(s):                            2\nNUMA node0 CPU(s):                       0-47,96-143\nNUMA node1 CPU(s):                       48-95,144-191\nVulnerability Gather data sampling:      Not affected\nVulnerability Indirect target selection: Not affected\nVulnerability Itlb multihit:             Not affected\nVulnerability L1tf:                      Not affected\nVulnerability Mds:                       Not affected\nVulnerability Meltdown:                  Not affected\nVulnerability Mmio stale data:           Not affected\nVulnerability Reg file data sampling:    Not affected\nVulnerability Retbleed:                  Not affected\nVulnerability Spec rstack overflow:      Not affected\nVulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl\nVulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization\nVulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S\nVulnerability Srbds:                     Not affected\nVulnerability Tsa:                       Not affected\nVulnerability Tsx async abort:           Not affected\nVulnerability Vmscape:                   Not affected"
+}

job.py

@@ -0,0 +1,394 @@
+# /// script
+# requires-python = "==3.10.*"
+# dependencies = [
+#   "spaces==0.50.1",
+#   "torch==2.12.0",
+#   "torchvision",
+#   "diffusers @ https://github.com/huggingface/diffusers/archive/refs/heads/main.tar.gz",
+#   "transformers",
+#   "accelerate",
+#   "sentencepiece",
+#   "imageio",
+#   "imageio-ffmpeg",
+#   "av",
+#   "safetensors",
+#   "ftfy",
+#   "numpy",
+#   "pillow",
+#   "huggingface_hub",
+#   "setuptools",
+# ]
+# ///
+
+# =========================
+# User section  (LTX-2.3 IC-LoRA, Group A: in-context AV, no self-attn mask, no STG)
+# =========================
+
+# README::MODEL_INIT::START
+import os
+import tempfile
+
+import numpy as np
+import torch
+import spaces
+from PIL import Image
+
+from diffusers import LTX2InContextPipeline
+from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
+from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
+
+# base == distilled in architecture, so one compiled graph serves both; distilled is
+# what most demos use. The AOTI package is weight-agnostic, so this base graph also
+# serves any FUSED LoRA (fuse_lora before aoti_load on the Space).
+MODEL_ID = os.environ.get("LTX_MODEL_ID", "diffusers/LTX-2.3-Distilled-Diffusers")
+
+pipe = LTX2InContextPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe.vae.enable_tiling()
+# README::MODEL_INIT::END
+
+FPS = 24
+WIDTH = int(os.environ.get("LTX_W", "768"))
+HEIGHT = int(os.environ.get("LTX_H", "448"))
+NUM_FRAMES = int(os.environ.get("LTX_FRAMES", "49"))
+NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
+SAMPLE_MODE = os.environ.get("LTX_SAMPLES", "stub")  # "stub" (cheap) or "real"
+# Group B: force the in-context self-attention mask. Setting conditioning_attention_strength
+# < 1.0 makes the pipeline build video_self_attention_mask (shape (B, T_v, T_v)) internally,
+# i.e. the same block-level self_attention_mask tensor inpaint/outpaint produce via a
+# pixel-space conditioning_attention_mask. Group A (default) leaves it None.
+GROUP_B = os.environ.get("LTX_GROUP_B", "0").strip().lower() in ("1", "true", "yes")
+COND_ATTN_STRENGTH = 0.9 if GROUP_B else 1.0
+# Group C (STG): compile the block with the perturbation path ALWAYS active so ONE graph
+# serves both STG and non-STG. We capture with perturbation off (perturbation_mask=None),
+# then inject a perturbation_mask TENSOR before export — the attention then traces
+# `lerp(value, hidden_states, perturbation_mask)`, which is a no-op when the mask is ones
+# (non-STG / non-STG blocks) and blends when it's the real STG mask (block 28, perturbed pass).
+# all_perturbed stays False so the `if all_perturbed:` python shortcut is dropped (equivalent:
+# lerp(value, hs, 0) == value). The Space passes ones when off / the real mask for the STG block.
+STG = os.environ.get("LTX_STG", "0").strip().lower() in ("1", "true", "yes")
+
+
+def _ref_frames(n, w, h):
+    """Synthetic grayscale reference frames (compilation needs valid shapes, not nice pixels)."""
+    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
+    out = []
+    for t in range(n):
+        g = (np.sin((xx / w + yy / h + t / max(n, 1)) * 2 * np.pi) * 0.5 + 0.5) * 255
+        out.append(Image.fromarray(g.astype(np.uint8)).convert("RGB"))
+    return out
+
+
+def _run_pipe(steps, output_path=None):
+    ref = _ref_frames(NUM_FRAMES, WIDTH, HEIGHT)
+    out = pipe(
+        prompt="a colorful natural scene with gentle ambient sound",
+        negative_prompt="",
+        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
+        reference_downscale_factor=1,
+        conditioning_attention_strength=COND_ATTN_STRENGTH,
+        width=WIDTH, height=HEIGHT, num_frames=NUM_FRAMES, frame_rate=FPS,
+        num_inference_steps=steps, sigmas=DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
+        generator=torch.Generator(device="cuda").manual_seed(0),
+        output_type="np", return_dict=False,
+    )
+    if output_path is not None:
+        from diffusers.utils import encode_video
+        video_np, audio = out[0], out[1]
+        kw = {}
+        if audio is not None:
+            kw = dict(audio=audio[0].float().cpu(),
+                      audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
+        encode_video(video_np[0], fps=FPS, output_path=output_path, **kw)
+    return out
+
+
+def _build_dynamic_shapes(block, call):
+    """Flat dynamic_shapes dict. The block forward has NO **kwargs (clean diffusers
+    signature), so no clean-forward hack is needed. Only the video-token count T_v and
+    audio-token count T_a vary (text is padded to a fixed 1024 -> static). Both are large,
+    so size-matching is collision-safe vs structural sizes (head_dim 128, heads 32,
+    caption 3840, etc.). Recurse into tuples (rotary embeddings are (cos, sin) pairs)."""
+    import inspect
+    from torch.export import Dim
+
+    posnames = [n for n, p in inspect.signature(type(block).forward).parameters.items()
+                if n != "self" and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)]
+    named = {posnames[i]: a for i, a in enumerate(call.args)}
+    named.update(call.kwargs or {})
+
+    T_v = named["hidden_states"].shape[1]
+    T_a = named["audio_hidden_states"].shape[1]
+    DYN = {T_v, T_a}
+
+    def spec(v):
+        if torch.is_tensor(v):
+            d = {i: Dim.DYNAMIC for i, s in enumerate(v.shape) if s in DYN}
+            return d or None
+        if isinstance(v, (list, tuple)):
+            return type(v)(spec(x) for x in v)
+        return None
+
+    return {k: spec(v) for k, v in named.items()}, T_v, T_a
+
+
+def compile_and_save(module: torch.nn.Module, package_dir: str):
+    submodule = "transformer_blocks"
+    block = module.get_submodule(submodule)[0]
+    with spaces.aoti_capture(block) as call:
+        _run_pipe(steps=NUM_STEPS)  # aoti_capture raises at the first block call
+
+    print("AOTI: captured block forward "
+          f"(args={len(call.args or ())}, kwargs={sorted((call.kwargs or {}).keys())})")
+    for k, v in (call.kwargs or {}).items():
+        if torch.is_tensor(v):
+            print(f"  {k}: Tensor {tuple(v.shape)} {v.dtype}")
+        elif isinstance(v, (list, tuple)):
+            inner = [tuple(x.shape) if torch.is_tensor(x) else type(x).__name__ for x in v]
+            print(f"  {k}: {type(v).__name__} {inner}")
+        else:
+            print(f"  {k}: {type(v).__name__} {v}")
+
+    if STG:
+        hs = call.kwargs.get("hidden_states", call.args[0] if call.args else None)
+        pm = torch.ones((hs.shape[0], 1, 1), device=hs.device, dtype=hs.dtype)
+        call.kwargs["perturbation_mask"] = pm       # None -> tensor: traces the lerp path
+        call.kwargs["all_perturbed"] = False         # drop the skip-attention python shortcut
+        print(f"AOTI: STG mode — injected perturbation_mask {tuple(pm.shape)}")
+
+    dynamic_shapes, T_v, T_a = _build_dynamic_shapes(block, call)
+    def _fmt(v):
+        if isinstance(v, dict):
+            return sorted(v)
+        if isinstance(v, (list, tuple)):
+            return [_fmt(x) for x in v]
+        return v
+    print(f"AOTI: T_v={T_v} T_a={T_a}; dynamic dims="
+          f"{ {k: _fmt(v) for k, v in dynamic_shapes.items() if v} }")
+
+    with torch.no_grad():
+        exported = torch.export.export(
+            block, args=call.args, kwargs=call.kwargs, dynamic_shapes=dynamic_shapes,
+        )
+    print("AOTI: torch.export OK")
+    spaces.aoti_compile_and_save(
+        package_dir=package_dir, exported_program=exported, submodule=submodule,
+    )
+    print("AOTI: compile_and_save OK")
+
+
+def generate_samples(samples_dir: str):
+    if SAMPLE_MODE != "real":
+        import imageio.v2 as imageio
+        frames = [(np.random.default_rng(i).random((64, 64, 3)) * 255).astype(np.uint8) for i in range(8)]
+        imageio.mimsave(f"{samples_dir}/video.mp4", frames, fps=8, macro_block_size=1)
+        return
+    _run_pipe(steps=NUM_STEPS, output_path=f"{samples_dir}/video.mp4")
+
+
+def main():
+    create_aoti_repo(
+        module=pipe.transformer,
+        module_expr="pipe.transformer",
+        compile_and_save=compile_and_save,
+        generate_samples=generate_samples,
+    )
+
+
+# =========================
+# Internal (avoid editing) — same harness as the reference AOTI job
+# =========================
+
+
+import inspect
+import json
+import random
+import shutil
+import sys
+import time
+from packaging.version import Version
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Callable
+
+import huggingface_hub as hf
+from requests.exceptions import HTTPError
+
+
+def create_aoti_repo(module, module_expr, compile_and_save, generate_samples, aoti_loader=None):
+    HUB_URL = 'https://huggingface.co'
+    user = hf.whoami()['name']
+    job_id = os.environ.get('JOB_ID')
+    job_info = hf.inspect_job(job_id=job_id) if job_id is not None else None
+    env_info = torch.utils.collect_env.get_env_info()
+    library_name, config = _get_library_config(module)
+
+    with TemporaryDirectory() as tempdir:
+        tempdir = Path(tempdir)
+        readme_path = tempdir / 'README.md'
+        package_dir = tempdir / 'package'
+        samples_before_dir = tempdir / 'samples' / 'before'
+        samples_after_dir = tempdir / 'samples' / 'after'
+        environment_path = tempdir / 'environment.json'
+        config_path = tempdir / 'module_config.json'
+
+        samples_before_dir.mkdir(parents=True)
+        t0 = time.perf_counter()
+        generate_samples(str(samples_before_dir))
+        generate_before_dt = time.perf_counter() - t0
+
+        package_dir.mkdir(parents=True)
+        compile_and_save(module, str(package_dir))
+        if aoti_loader is not None:
+            aoti_loader(module, str(package_dir))
+        else:
+            spaces.aoti_load_from_package_dir(module, package_dir)
+
+        samples_after_dir.mkdir(parents=True)
+        t0 = time.perf_counter()
+        generate_samples(str(samples_after_dir))
+        generate_after_dt = time.perf_counter() - t0
+
+        environment_path.write_text(json.dumps(env_info._asdict(), indent=4))
+        if config is not None:
+            config_path.write_text(json.dumps(config, indent=4))
+
+        output_repo_id = _create_empty_repo(
+            user=user, module=module, cuda_version=env_info.cuda_runtime_version,
+            kernels=(package_dir / 'kernels').is_dir(),
+        )
+
+        model_init_region = (inspect.getsource(sys.modules['__main__'])
+            .split('\n# README::MODEL_INIT::START')[1]
+            .split('\n# README::MODEL_INIT::END')[0])
+        aoti_load_readme = spaces.aoti_load_call_source(
+            module_expr=module_expr, repo_id=output_repo_id, aoti_loader=aoti_loader)
+        def get_link(path: Path):
+            kind = 'tree' if path.is_dir() else 'resolve'
+            return f'{HUB_URL}/{output_repo_id}/{kind}/main/{path.relative_to(tempdir)}'
+        readme_path.write_text(_readme_template(
+            model_init=model_init_region, aoti_load=aoti_load_readme, repo_id=output_repo_id,
+            job_id=f'{user}/{job_id}',
+            job_image=job_info.docker_image if job_info is not None else os.getenv('JOB_IMAGE'),
+            job_flavor=job_info.flavor if job_info is not None else os.getenv('JOB_FLAVOR'),
+            environment=torch.utils.collect_env.pretty_str(env_info),
+            library_name=library_name,
+            generate_before_dt=generate_before_dt, generate_after_dt=generate_after_dt,
+            samples_before_urls=[get_link(p) for p in samples_before_dir.iterdir()],
+            samples_after_urls=[get_link(p) for p in samples_after_dir.iterdir()],
+        ))
+        shutil.copyfile(__file__, tempdir / 'job.py')
+        hf.upload_folder(repo_id=output_repo_id, folder_path=tempdir)
+        print(f"AoT repository successfully created at: {HUB_URL}/{output_repo_id}")
+
+
+def _create_empty_repo(user, module, cuda_version, kernels, max_attempts=10):
+    for _ in range(max_attempts):
+        output_repo_id = _get_repo_id(user, module, cuda_version, kernels)
+        try:
+            hf.create_repo(output_repo_id, private=True)
+        except HTTPError as err:
+            if err.response.status_code != 409:
+                raise
+        else:
+            return output_repo_id
+    raise AssertionError
+
+
+def _get_repo_id(user, module, cuda_version, kernels):
+    if (repo_id := os.getenv('OUTPUT_REPO_ID')) is not None:
+        return repo_id
+    namespace = os.getenv('OUTPUT_REPO_NAMESPACE', user)
+    base_name = os.getenv('OUTPUT_REPO_BASE_NAME', module.__class__.__name__)
+    sm = ''.join(map(str, torch.cuda.get_device_capability()))
+    cu = ''.join(cuda_version.split('.')[:2])
+    rnd = random.randbytes(1).hex()
+    res = f'{namespace}/{base_name}-sm{sm}-cu{cu}'
+    if kernels:
+        torch_version = Version(torch.__version__)
+        res += f'-torch{torch_version.major}{torch_version.minor}'
+    return f'{res}-r{rnd}'
+
+
+def _get_library_config(module):
+    if (config := getattr(module, 'config', None)) is None:
+        return None, None
+    if callable(getattr(config, 'to_dict', None)):
+        config = config.to_dict()
+    if not isinstance(config, dict):
+        return None, None
+    if 'transformers_version' in config:
+        library_name = 'transformers'
+    elif '_diffusers_version' in config:
+        library_name = 'diffusers'
+    else:
+        library_name = 'unknown'
+    return library_name, config
+
+
+def _readme_template(model_init, aoti_load, repo_id, job_id, job_image, job_flavor,
+                     environment, library_name, generate_before_dt, generate_after_dt,
+                     samples_before_urls, samples_after_urls):
+    NEWLINE = '\n'
+    IMAGE_EXTS = ('.png', '.webp', '.jpg', '.jpeg', '.gif')
+    VIDEO_EXTS = ('.mp4', '.webm', '.mov')
+    def media_cell(url):
+        name = url.split('/')[-1]
+        if name.endswith(IMAGE_EXTS):
+            return f'![{name}]({url})'
+        if name.endswith(VIDEO_EXTS):
+            return f'<video src="{url}" controls></video>'
+        return f'[{name}]({url})'
+    return f"""
+---
+tags:
+- ahead-of-time
+- pytorch
+library_name: {library_name or 'pytorch'}
+---
+
+> [!NOTE]
+> This **README** has been auto-generated by the **HF Job** run linked below
+> and the whole repository is a reproducible artifact of this Job
+
+# Ahead-of-time repository
+
+AoT repos contain **pre-compiled binaries** of PyTorch models, enabling:
+- fast startup times (no `torch.compile` needed)
+- significant **speedup**
+- **ZeroGPU** compatibility
+
+## How to use
+``` python
+{model_init}\n
+{aoti_load}
+```
+
+## How to reproduce or customize
+``` bash
+hf jobs uv run job.py --flavor {job_flavor or '<unknown>'} --image {job_image or '<unknown>'} --secrets HF_TOKEN
+```
+
+## Samples
+| Before compilation ({generate_before_dt:.2f}s) | After compilation ({generate_after_dt:.2f}s) |
+|---|---|
+{NEWLINE.join(f"| {media_cell(b)} | {media_cell(a)} |" for b, a in zip(samples_before_urls, samples_after_urls))}
+
+Speedup: **{generate_before_dt/generate_after_dt:.2f}x**
+
+## Environment
+<details><summary>Click to expand</summary>
+
+```
+{environment}
+```
+</details>
+
+## Job run
+- [{job_id}](https://huggingface.co/jobs/{job_id})
+"""
+
+
+if __name__ == '__main__':
+    main()

module_config.json

@@ -0,0 +1,51 @@
+{
+    "in_channels": 128,
+    "out_channels": 128,
+    "patch_size": 1,
+    "patch_size_t": 1,
+    "num_attention_heads": 32,
+    "attention_head_dim": 128,
+    "cross_attention_dim": 4096,
+    "vae_scale_factors": [
+        8,
+        32,
+        32
+    ],
+    "pos_embed_max_pos": 20,
+    "base_height": 2048,
+    "base_width": 2048,
+    "gated_attn": true,
+    "cross_attn_mod": true,
+    "audio_in_channels": 128,
+    "audio_out_channels": 128,
+    "audio_patch_size": 1,
+    "audio_patch_size_t": 1,
+    "audio_num_attention_heads": 32,
+    "audio_attention_head_dim": 64,
+    "audio_cross_attention_dim": 2048,
+    "audio_scale_factor": 4,
+    "audio_pos_embed_max_pos": 20,
+    "audio_sampling_rate": 16000,
+    "audio_hop_length": 160,
+    "audio_gated_attn": true,
+    "audio_cross_attn_mod": true,
+    "num_layers": 48,
+    "activation_fn": "gelu-approximate",
+    "qk_norm": "rms_norm_across_heads",
+    "norm_elementwise_affine": false,
+    "norm_eps": 1e-06,
+    "caption_channels": 3840,
+    "attention_bias": true,
+    "attention_out_bias": true,
+    "rope_theta": 10000.0,
+    "rope_double_precision": true,
+    "causal_offset": 1,
+    "timestep_scale_multiplier": 1000,
+    "cross_attn_timestep_scale_multiplier": 1000,
+    "rope_type": "split",
+    "use_prompt_embeddings": false,
+    "perturbed_attn": true,
+    "_class_name": "LTX2VideoTransformer3DModel",
+    "_diffusers_version": "0.37.0.dev0",
+    "_name_or_path": "/root/.cache/huggingface/hub/models--diffusers--LTX-2.3-Distilled-Diffusers/snapshots/432e0d3c2d1769aaa4d295f9243f7062bf6b47ee/transformer"
+}

package/submodules/transformer_blocks/package.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50ee44ed94340ab0491f2c8909d3746f2683eb816b6c7b3bac5a849a2c141a8c
+size 2139804