linoyts HF Staff commited on
Commit
ab90d75
·
verified ·
1 Parent(s): 43a4afd

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ package/submodules/transformer_blocks/package.pt2 filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,161 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ ---
3
+ tags:
4
+ - ahead-of-time
5
+ - pytorch
6
+ library_name: diffusers
7
+ ---
8
+
9
+ > [!NOTE]
10
+ > This **README** has been auto-generated by the **HF Job** run linked below
11
+ > and the whole repository is a reproducible artifact of this Job
12
+
13
+ # Ahead-of-time repository
14
+
15
+ AoT repos contain **pre-compiled binaries** of PyTorch models, enabling:
16
+ - fast startup times (no `torch.compile` needed)
17
+ - significant **speedup**
18
+ - **ZeroGPU** compatibility
19
+
20
+ ## How to use
21
+ ``` python
22
+
23
+ import os
24
+ import tempfile
25
+
26
+ import numpy as np
27
+ import torch
28
+ import spaces
29
+ from PIL import Image
30
+
31
+ from diffusers import LTX2InContextPipeline
32
+ from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
33
+ from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
34
+
35
+ # base == distilled in architecture, so one compiled graph serves both; distilled is
36
+ # what most demos use. The AOTI package is weight-agnostic, so this base graph also
37
+ # serves any FUSED LoRA (fuse_lora before aoti_load on the Space).
38
+ MODEL_ID = os.environ.get("LTX_MODEL_ID", "diffusers/LTX-2.3-Distilled-Diffusers")
39
+
40
+ pipe = LTX2InContextPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
41
+ pipe.to("cuda")
42
+ pipe.vae.enable_tiling()
43
+
44
+ spaces.aoti_load(
45
+ module=pipe.transformer,
46
+ repo_id='linoyts/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e',
47
+ )
48
+ ```
49
+
50
+ ## How to reproduce or customize
51
+ ``` bash
52
+ hf jobs uv run job.py --flavor rtx-pro-6000 --image pytorch/pytorch:2.9.1-cuda13.0-cudnn9-devel --secrets HF_TOKEN
53
+ ```
54
+
55
+ ## Samples
56
+ | Before compilation (0.37s) | After compilation (0.30s) |
57
+ |---|---|
58
+ | <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e/resolve/main/samples/before/video.mp4" controls></video> | <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e/resolve/main/samples/after/video.mp4" controls></video> |
59
+
60
+ Speedup: **1.21x**
61
+
62
+ ## Environment
63
+ <details><summary>Click to expand</summary>
64
+
65
+ ```
66
+ PyTorch version: 2.12.0+cu130
67
+ Is debug build: False
68
+ CUDA used to build PyTorch: 13.0
69
+ ROCM used to build PyTorch: N/A
70
+
71
+ OS: Ubuntu 22.04.5 LTS (x86_64)
72
+ GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
73
+ Clang version: Could not collect
74
+ CMake version: version 4.1.2
75
+ Libc version: glibc-2.35
76
+
77
+ Python version: 3.10.19 (main, Oct 31 2025, 23:02:46) [Clang 21.1.4 ] (64-bit runtime)
78
+ Python platform: Linux-6.12.88-119.157.amzn2023.x86_64-x86_64-with-glibc2.35
79
+ Is CUDA available: True
80
+ CUDA runtime version: 13.0.48
81
+ CUDA_MODULE_LOADING set to:
82
+ GPU models and configuration: GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition
83
+ Nvidia driver version: 580.159.03
84
+ cuDNN version: Could not collect
85
+ Is XPU available: False
86
+ HIP runtime version: N/A
87
+ MIOpen runtime version: N/A
88
+ Is XNNPACK available: True
89
+ Caching allocator config: N/A
90
+
91
+ CPU:
92
+ Architecture: x86_64
93
+ CPU op-mode(s): 32-bit, 64-bit
94
+ Address sizes: 46 bits physical, 48 bits virtual
95
+ Byte Order: Little Endian
96
+ CPU(s): 192
97
+ On-line CPU(s) list: 0-191
98
+ Vendor ID: GenuineIntel
99
+ Model name: Intel(R) Xeon(R) Platinum 8559C
100
+ CPU family: 6
101
+ Model: 207
102
+ Thread(s) per core: 2
103
+ Core(s) per socket: 48
104
+ Socket(s): 2
105
+ Stepping: 2
106
+ BogoMIPS: 4800.00
107
+ 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
108
+ Hypervisor vendor: KVM
109
+ Virtualization type: full
110
+ L1d cache: 4.5 MiB (96 instances)
111
+ L1i cache: 3 MiB (96 instances)
112
+ L2 cache: 192 MiB (96 instances)
113
+ L3 cache: 640 MiB (2 instances)
114
+ NUMA node(s): 2
115
+ NUMA node0 CPU(s): 0-47,96-143
116
+ NUMA node1 CPU(s): 48-95,144-191
117
+ Vulnerability Gather data sampling: Not affected
118
+ Vulnerability Indirect target selection: Not affected
119
+ Vulnerability Itlb multihit: Not affected
120
+ Vulnerability L1tf: Not affected
121
+ Vulnerability Mds: Not affected
122
+ Vulnerability Meltdown: Not affected
123
+ Vulnerability Mmio stale data: Not affected
124
+ Vulnerability Reg file data sampling: Not affected
125
+ Vulnerability Retbleed: Not affected
126
+ Vulnerability Spec rstack overflow: Not affected
127
+ Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
128
+ Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
129
+ Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
130
+ Vulnerability Srbds: Not affected
131
+ Vulnerability Tsa: Not affected
132
+ Vulnerability Tsx async abort: Not affected
133
+ Vulnerability Vmscape: Not affected
134
+
135
+ Versions of relevant libraries:
136
+ [pip3] Could not collect
137
+ [conda] numpy 2.3.4 py311h2e04523_0 conda-forge
138
+ [conda] nvidia-cublas 13.0.0.19 pypi_0 pypi
139
+ [conda] nvidia-cuda-cupti 13.0.48 pypi_0 pypi
140
+ [conda] nvidia-cuda-nvrtc 13.0.48 pypi_0 pypi
141
+ [conda] nvidia-cuda-runtime 13.0.48 pypi_0 pypi
142
+ [conda] nvidia-cudnn-cu13 9.13.0.50 pypi_0 pypi
143
+ [conda] nvidia-cufft 12.0.0.15 pypi_0 pypi
144
+ [conda] nvidia-curand 10.4.0.35 pypi_0 pypi
145
+ [conda] nvidia-cusolver 12.0.3.29 pypi_0 pypi
146
+ [conda] nvidia-cusparse 12.6.2.49 pypi_0 pypi
147
+ [conda] nvidia-cusparselt-cu13 0.8.0 pypi_0 pypi
148
+ [conda] nvidia-nccl-cu13 2.27.7 pypi_0 pypi
149
+ [conda] nvidia-nvjitlink 13.0.39 pypi_0 pypi
150
+ [conda] nvidia-nvtx 13.0.39 pypi_0 pypi
151
+ [conda] optree 0.17.0 pypi_0 pypi
152
+ [conda] torch 2.9.1+cu130 pypi_0 pypi
153
+ [conda] torchaudio 2.9.1+cu130 pypi_0 pypi
154
+ [conda] torchelastic 0.2.2 pypi_0 pypi
155
+ [conda] torchvision 0.24.1+cu130 pypi_0 pypi
156
+ [conda] triton 3.5.1 pypi_0 pypi
157
+ ```
158
+ </details>
159
+
160
+ ## Job run
161
+ - [linoyts/6a326f7f5ff0a6cf94fa00e8](https://huggingface.co/jobs/linoyts/6a326f7f5ff0a6cf94fa00e8)
environment.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "torch_version": "2.12.0+cu130",
3
+ "is_debug_build": "False",
4
+ "cuda_compiled_version": "13.0",
5
+ "gcc_version": "(Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0",
6
+ "clang_version": null,
7
+ "cmake_version": "version 4.1.2",
8
+ "os": "Ubuntu 22.04.5 LTS (x86_64)",
9
+ "libc_version": "glibc-2.35",
10
+ "python_version": "3.10.19 (main, Oct 31 2025, 23:02:46) [Clang 21.1.4 ] (64-bit runtime)",
11
+ "python_platform": "Linux-6.12.88-119.157.amzn2023.x86_64-x86_64-with-glibc2.35",
12
+ "is_cuda_available": "True",
13
+ "cuda_runtime_version": "13.0.48",
14
+ "cuda_module_loading": "",
15
+ "nvidia_driver_version": "580.159.03",
16
+ "nvidia_gpu_models": "GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition",
17
+ "cudnn_version": null,
18
+ "is_xpu_available": "False",
19
+ "pip_version": "pip3",
20
+ "pip_packages": null,
21
+ "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",
22
+ "hip_compiled_version": "N/A",
23
+ "hip_runtime_version": "N/A",
24
+ "miopen_runtime_version": "N/A",
25
+ "caching_allocator_config": {},
26
+ "is_xnnpack_available": "True",
27
+ "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"
28
+ }
job.py ADDED
@@ -0,0 +1,372 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # requires-python = "==3.10.*"
3
+ # dependencies = [
4
+ # "spaces==0.50.1",
5
+ # "torch==2.12.0",
6
+ # "torchvision",
7
+ # "diffusers @ https://github.com/huggingface/diffusers/archive/refs/heads/main.tar.gz",
8
+ # "transformers",
9
+ # "accelerate",
10
+ # "sentencepiece",
11
+ # "imageio",
12
+ # "imageio-ffmpeg",
13
+ # "av",
14
+ # "safetensors",
15
+ # "ftfy",
16
+ # "numpy",
17
+ # "pillow",
18
+ # "huggingface_hub",
19
+ # "setuptools",
20
+ # ]
21
+ # ///
22
+
23
+ # =========================
24
+ # User section (LTX-2.3 IC-LoRA, Group A: in-context AV, no self-attn mask, no STG)
25
+ # =========================
26
+
27
+ # README::MODEL_INIT::START
28
+ import os
29
+ import tempfile
30
+
31
+ import numpy as np
32
+ import torch
33
+ import spaces
34
+ from PIL import Image
35
+
36
+ from diffusers import LTX2InContextPipeline
37
+ from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
38
+ from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
39
+
40
+ # base == distilled in architecture, so one compiled graph serves both; distilled is
41
+ # what most demos use. The AOTI package is weight-agnostic, so this base graph also
42
+ # serves any FUSED LoRA (fuse_lora before aoti_load on the Space).
43
+ MODEL_ID = os.environ.get("LTX_MODEL_ID", "diffusers/LTX-2.3-Distilled-Diffusers")
44
+
45
+ pipe = LTX2InContextPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
46
+ pipe.to("cuda")
47
+ pipe.vae.enable_tiling()
48
+ # README::MODEL_INIT::END
49
+
50
+ FPS = 24
51
+ WIDTH = int(os.environ.get("LTX_W", "768"))
52
+ HEIGHT = int(os.environ.get("LTX_H", "448"))
53
+ NUM_FRAMES = int(os.environ.get("LTX_FRAMES", "49"))
54
+ NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
55
+ SAMPLE_MODE = os.environ.get("LTX_SAMPLES", "stub") # "stub" (cheap) or "real"
56
+
57
+
58
+ def _ref_frames(n, w, h):
59
+ """Synthetic grayscale reference frames (compilation needs valid shapes, not nice pixels)."""
60
+ yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
61
+ out = []
62
+ for t in range(n):
63
+ g = (np.sin((xx / w + yy / h + t / max(n, 1)) * 2 * np.pi) * 0.5 + 0.5) * 255
64
+ out.append(Image.fromarray(g.astype(np.uint8)).convert("RGB"))
65
+ return out
66
+
67
+
68
+ def _run_pipe(steps, output_path=None):
69
+ ref = _ref_frames(NUM_FRAMES, WIDTH, HEIGHT)
70
+ out = pipe(
71
+ prompt="a colorful natural scene with gentle ambient sound",
72
+ negative_prompt="",
73
+ reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
74
+ reference_downscale_factor=1,
75
+ width=WIDTH, height=HEIGHT, num_frames=NUM_FRAMES, frame_rate=FPS,
76
+ num_inference_steps=steps, sigmas=DISTILLED_SIGMA_VALUES,
77
+ guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
78
+ generator=torch.Generator(device="cuda").manual_seed(0),
79
+ output_type="np", return_dict=False,
80
+ )
81
+ if output_path is not None:
82
+ from diffusers.utils import encode_video
83
+ video_np, audio = out[0], out[1]
84
+ kw = {}
85
+ if audio is not None:
86
+ kw = dict(audio=audio[0].float().cpu(),
87
+ audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
88
+ encode_video(video_np[0], fps=FPS, output_path=output_path, **kw)
89
+ return out
90
+
91
+
92
+ def _build_dynamic_shapes(block, call):
93
+ """Flat dynamic_shapes dict. The block forward has NO **kwargs (clean diffusers
94
+ signature), so no clean-forward hack is needed. Only the video-token count T_v and
95
+ audio-token count T_a vary (text is padded to a fixed 1024 -> static). Both are large,
96
+ so size-matching is collision-safe vs structural sizes (head_dim 128, heads 32,
97
+ caption 3840, etc.). Recurse into tuples (rotary embeddings are (cos, sin) pairs)."""
98
+ import inspect
99
+ from torch.export import Dim
100
+
101
+ posnames = [n for n, p in inspect.signature(type(block).forward).parameters.items()
102
+ if n != "self" and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)]
103
+ named = {posnames[i]: a for i, a in enumerate(call.args)}
104
+ named.update(call.kwargs or {})
105
+
106
+ T_v = named["hidden_states"].shape[1]
107
+ T_a = named["audio_hidden_states"].shape[1]
108
+ DYN = {T_v, T_a}
109
+
110
+ def spec(v):
111
+ if torch.is_tensor(v):
112
+ d = {i: Dim.DYNAMIC for i, s in enumerate(v.shape) if s in DYN}
113
+ return d or None
114
+ if isinstance(v, (list, tuple)):
115
+ return type(v)(spec(x) for x in v)
116
+ return None
117
+
118
+ return {k: spec(v) for k, v in named.items()}, T_v, T_a
119
+
120
+
121
+ def compile_and_save(module: torch.nn.Module, package_dir: str):
122
+ submodule = "transformer_blocks"
123
+ block = module.get_submodule(submodule)[0]
124
+ with spaces.aoti_capture(block) as call:
125
+ _run_pipe(steps=NUM_STEPS) # aoti_capture raises at the first block call
126
+
127
+ print("AOTI: captured block forward "
128
+ f"(args={len(call.args or ())}, kwargs={sorted((call.kwargs or {}).keys())})")
129
+ for k, v in (call.kwargs or {}).items():
130
+ if torch.is_tensor(v):
131
+ print(f" {k}: Tensor {tuple(v.shape)} {v.dtype}")
132
+ elif isinstance(v, (list, tuple)):
133
+ inner = [tuple(x.shape) if torch.is_tensor(x) else type(x).__name__ for x in v]
134
+ print(f" {k}: {type(v).__name__} {inner}")
135
+ else:
136
+ print(f" {k}: {type(v).__name__} {v}")
137
+
138
+ dynamic_shapes, T_v, T_a = _build_dynamic_shapes(block, call)
139
+ def _fmt(v):
140
+ if isinstance(v, dict):
141
+ return sorted(v)
142
+ if isinstance(v, (list, tuple)):
143
+ return [_fmt(x) for x in v]
144
+ return v
145
+ print(f"AOTI: T_v={T_v} T_a={T_a}; dynamic dims="
146
+ f"{ {k: _fmt(v) for k, v in dynamic_shapes.items() if v} }")
147
+
148
+ with torch.no_grad():
149
+ exported = torch.export.export(
150
+ block, args=call.args, kwargs=call.kwargs, dynamic_shapes=dynamic_shapes,
151
+ )
152
+ print("AOTI: torch.export OK")
153
+ spaces.aoti_compile_and_save(
154
+ package_dir=package_dir, exported_program=exported, submodule=submodule,
155
+ )
156
+ print("AOTI: compile_and_save OK")
157
+
158
+
159
+ def generate_samples(samples_dir: str):
160
+ if SAMPLE_MODE != "real":
161
+ import imageio.v2 as imageio
162
+ frames = [(np.random.default_rng(i).random((64, 64, 3)) * 255).astype(np.uint8) for i in range(8)]
163
+ imageio.mimsave(f"{samples_dir}/video.mp4", frames, fps=8, macro_block_size=1)
164
+ return
165
+ _run_pipe(steps=NUM_STEPS, output_path=f"{samples_dir}/video.mp4")
166
+
167
+
168
+ def main():
169
+ create_aoti_repo(
170
+ module=pipe.transformer,
171
+ module_expr="pipe.transformer",
172
+ compile_and_save=compile_and_save,
173
+ generate_samples=generate_samples,
174
+ )
175
+
176
+
177
+ # =========================
178
+ # Internal (avoid editing) — same harness as the reference AOTI job
179
+ # =========================
180
+
181
+
182
+ import inspect
183
+ import json
184
+ import random
185
+ import shutil
186
+ import sys
187
+ import time
188
+ from packaging.version import Version
189
+ from pathlib import Path
190
+ from tempfile import TemporaryDirectory
191
+ from typing import Callable
192
+
193
+ import huggingface_hub as hf
194
+ from requests.exceptions import HTTPError
195
+
196
+
197
+ def create_aoti_repo(module, module_expr, compile_and_save, generate_samples, aoti_loader=None):
198
+ HUB_URL = 'https://huggingface.co'
199
+ user = hf.whoami()['name']
200
+ job_id = os.environ.get('JOB_ID')
201
+ job_info = hf.inspect_job(job_id=job_id) if job_id is not None else None
202
+ env_info = torch.utils.collect_env.get_env_info()
203
+ library_name, config = _get_library_config(module)
204
+
205
+ with TemporaryDirectory() as tempdir:
206
+ tempdir = Path(tempdir)
207
+ readme_path = tempdir / 'README.md'
208
+ package_dir = tempdir / 'package'
209
+ samples_before_dir = tempdir / 'samples' / 'before'
210
+ samples_after_dir = tempdir / 'samples' / 'after'
211
+ environment_path = tempdir / 'environment.json'
212
+ config_path = tempdir / 'module_config.json'
213
+
214
+ samples_before_dir.mkdir(parents=True)
215
+ t0 = time.perf_counter()
216
+ generate_samples(str(samples_before_dir))
217
+ generate_before_dt = time.perf_counter() - t0
218
+
219
+ package_dir.mkdir(parents=True)
220
+ compile_and_save(module, str(package_dir))
221
+ if aoti_loader is not None:
222
+ aoti_loader(module, str(package_dir))
223
+ else:
224
+ spaces.aoti_load_from_package_dir(module, package_dir)
225
+
226
+ samples_after_dir.mkdir(parents=True)
227
+ t0 = time.perf_counter()
228
+ generate_samples(str(samples_after_dir))
229
+ generate_after_dt = time.perf_counter() - t0
230
+
231
+ environment_path.write_text(json.dumps(env_info._asdict(), indent=4))
232
+ if config is not None:
233
+ config_path.write_text(json.dumps(config, indent=4))
234
+
235
+ output_repo_id = _create_empty_repo(
236
+ user=user, module=module, cuda_version=env_info.cuda_runtime_version,
237
+ kernels=(package_dir / 'kernels').is_dir(),
238
+ )
239
+
240
+ model_init_region = (inspect.getsource(sys.modules['__main__'])
241
+ .split('\n# README::MODEL_INIT::START')[1]
242
+ .split('\n# README::MODEL_INIT::END')[0])
243
+ aoti_load_readme = spaces.aoti_load_call_source(
244
+ module_expr=module_expr, repo_id=output_repo_id, aoti_loader=aoti_loader)
245
+ def get_link(path: Path):
246
+ kind = 'tree' if path.is_dir() else 'resolve'
247
+ return f'{HUB_URL}/{output_repo_id}/{kind}/main/{path.relative_to(tempdir)}'
248
+ readme_path.write_text(_readme_template(
249
+ model_init=model_init_region, aoti_load=aoti_load_readme, repo_id=output_repo_id,
250
+ job_id=f'{user}/{job_id}',
251
+ job_image=job_info.docker_image if job_info is not None else os.getenv('JOB_IMAGE'),
252
+ job_flavor=job_info.flavor if job_info is not None else os.getenv('JOB_FLAVOR'),
253
+ environment=torch.utils.collect_env.pretty_str(env_info),
254
+ library_name=library_name,
255
+ generate_before_dt=generate_before_dt, generate_after_dt=generate_after_dt,
256
+ samples_before_urls=[get_link(p) for p in samples_before_dir.iterdir()],
257
+ samples_after_urls=[get_link(p) for p in samples_after_dir.iterdir()],
258
+ ))
259
+ shutil.copyfile(__file__, tempdir / 'job.py')
260
+ hf.upload_folder(repo_id=output_repo_id, folder_path=tempdir)
261
+ print(f"AoT repository successfully created at: {HUB_URL}/{output_repo_id}")
262
+
263
+
264
+ def _create_empty_repo(user, module, cuda_version, kernels, max_attempts=10):
265
+ for _ in range(max_attempts):
266
+ output_repo_id = _get_repo_id(user, module, cuda_version, kernels)
267
+ try:
268
+ hf.create_repo(output_repo_id, private=True)
269
+ except HTTPError as err:
270
+ if err.response.status_code != 409:
271
+ raise
272
+ else:
273
+ return output_repo_id
274
+ raise AssertionError
275
+
276
+
277
+ def _get_repo_id(user, module, cuda_version, kernels):
278
+ if (repo_id := os.getenv('OUTPUT_REPO_ID')) is not None:
279
+ return repo_id
280
+ namespace = os.getenv('OUTPUT_REPO_NAMESPACE', user)
281
+ base_name = os.getenv('OUTPUT_REPO_BASE_NAME', module.__class__.__name__)
282
+ sm = ''.join(map(str, torch.cuda.get_device_capability()))
283
+ cu = ''.join(cuda_version.split('.')[:2])
284
+ rnd = random.randbytes(1).hex()
285
+ res = f'{namespace}/{base_name}-sm{sm}-cu{cu}'
286
+ if kernels:
287
+ torch_version = Version(torch.__version__)
288
+ res += f'-torch{torch_version.major}{torch_version.minor}'
289
+ return f'{res}-r{rnd}'
290
+
291
+
292
+ def _get_library_config(module):
293
+ if (config := getattr(module, 'config', None)) is None:
294
+ return None, None
295
+ if callable(getattr(config, 'to_dict', None)):
296
+ config = config.to_dict()
297
+ if not isinstance(config, dict):
298
+ return None, None
299
+ if 'transformers_version' in config:
300
+ library_name = 'transformers'
301
+ elif '_diffusers_version' in config:
302
+ library_name = 'diffusers'
303
+ else:
304
+ library_name = 'unknown'
305
+ return library_name, config
306
+
307
+
308
+ def _readme_template(model_init, aoti_load, repo_id, job_id, job_image, job_flavor,
309
+ environment, library_name, generate_before_dt, generate_after_dt,
310
+ samples_before_urls, samples_after_urls):
311
+ NEWLINE = '\n'
312
+ IMAGE_EXTS = ('.png', '.webp', '.jpg', '.jpeg', '.gif')
313
+ VIDEO_EXTS = ('.mp4', '.webm', '.mov')
314
+ def media_cell(url):
315
+ name = url.split('/')[-1]
316
+ if name.endswith(IMAGE_EXTS):
317
+ return f'![{name}]({url})'
318
+ if name.endswith(VIDEO_EXTS):
319
+ return f'<video src="{url}" controls></video>'
320
+ return f'[{name}]({url})'
321
+ return f"""
322
+ ---
323
+ tags:
324
+ - ahead-of-time
325
+ - pytorch
326
+ library_name: {library_name or 'pytorch'}
327
+ ---
328
+
329
+ > [!NOTE]
330
+ > This **README** has been auto-generated by the **HF Job** run linked below
331
+ > and the whole repository is a reproducible artifact of this Job
332
+
333
+ # Ahead-of-time repository
334
+
335
+ AoT repos contain **pre-compiled binaries** of PyTorch models, enabling:
336
+ - fast startup times (no `torch.compile` needed)
337
+ - significant **speedup**
338
+ - **ZeroGPU** compatibility
339
+
340
+ ## How to use
341
+ ``` python
342
+ {model_init}\n
343
+ {aoti_load}
344
+ ```
345
+
346
+ ## How to reproduce or customize
347
+ ``` bash
348
+ hf jobs uv run job.py --flavor {job_flavor or '<unknown>'} --image {job_image or '<unknown>'} --secrets HF_TOKEN
349
+ ```
350
+
351
+ ## Samples
352
+ | Before compilation ({generate_before_dt:.2f}s) | After compilation ({generate_after_dt:.2f}s) |
353
+ |---|---|
354
+ {NEWLINE.join(f"| {media_cell(b)} | {media_cell(a)} |" for b, a in zip(samples_before_urls, samples_after_urls))}
355
+
356
+ Speedup: **{generate_before_dt/generate_after_dt:.2f}x**
357
+
358
+ ## Environment
359
+ <details><summary>Click to expand</summary>
360
+
361
+ ```
362
+ {environment}
363
+ ```
364
+ </details>
365
+
366
+ ## Job run
367
+ - [{job_id}](https://huggingface.co/jobs/{job_id})
368
+ """
369
+
370
+
371
+ if __name__ == '__main__':
372
+ main()
module_config.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "in_channels": 128,
3
+ "out_channels": 128,
4
+ "patch_size": 1,
5
+ "patch_size_t": 1,
6
+ "num_attention_heads": 32,
7
+ "attention_head_dim": 128,
8
+ "cross_attention_dim": 4096,
9
+ "vae_scale_factors": [
10
+ 8,
11
+ 32,
12
+ 32
13
+ ],
14
+ "pos_embed_max_pos": 20,
15
+ "base_height": 2048,
16
+ "base_width": 2048,
17
+ "gated_attn": true,
18
+ "cross_attn_mod": true,
19
+ "audio_in_channels": 128,
20
+ "audio_out_channels": 128,
21
+ "audio_patch_size": 1,
22
+ "audio_patch_size_t": 1,
23
+ "audio_num_attention_heads": 32,
24
+ "audio_attention_head_dim": 64,
25
+ "audio_cross_attention_dim": 2048,
26
+ "audio_scale_factor": 4,
27
+ "audio_pos_embed_max_pos": 20,
28
+ "audio_sampling_rate": 16000,
29
+ "audio_hop_length": 160,
30
+ "audio_gated_attn": true,
31
+ "audio_cross_attn_mod": true,
32
+ "num_layers": 48,
33
+ "activation_fn": "gelu-approximate",
34
+ "qk_norm": "rms_norm_across_heads",
35
+ "norm_elementwise_affine": false,
36
+ "norm_eps": 1e-06,
37
+ "caption_channels": 3840,
38
+ "attention_bias": true,
39
+ "attention_out_bias": true,
40
+ "rope_theta": 10000.0,
41
+ "rope_double_precision": true,
42
+ "causal_offset": 1,
43
+ "timestep_scale_multiplier": 1000,
44
+ "cross_attn_timestep_scale_multiplier": 1000,
45
+ "rope_type": "split",
46
+ "use_prompt_embeddings": false,
47
+ "perturbed_attn": true,
48
+ "_class_name": "LTX2VideoTransformer3DModel",
49
+ "_diffusers_version": "0.37.0.dev0",
50
+ "_name_or_path": "/root/.cache/huggingface/hub/models--diffusers--LTX-2.3-Distilled-Diffusers/snapshots/432e0d3c2d1769aaa4d295f9243f7062bf6b47ee/transformer"
51
+ }
package/submodules/transformer_blocks/package.pt2 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e119689fc27e1f958ffdd68bb6405762edc045ce2508586cd761269505d31e48
3
+ size 2090702
samples/after/video.mp4 ADDED
Binary file (21.8 kB). View file
 
samples/before/video.mp4 ADDED
Binary file (21.8 kB). View file