linoyts HF Staff commited on
Commit
c3b6cd1
·
verified ·
1 Parent(s): 7f7cee5

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-GroupB-sm120-cu130-r0e',
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.36s) | After compilation (0.30s) |
57
+ |---|---|
58
+ | <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupB-sm120-cu130-r0e/resolve/main/samples/before/video.mp4" controls></video> | <video src="https://huggingface.co/linoyts/LTX-2.3-Transformer-GroupB-sm120-cu130-r0e/resolve/main/samples/after/video.mp4" controls></video> |
59
+
60
+ Speedup: **1.19x**
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/6a3270a3fb114ff24a388567](https://huggingface.co/jobs/linoyts/6a3270a3fb114ff24a388567)
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,379 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ # Group B: force the in-context self-attention mask. Setting conditioning_attention_strength
57
+ # < 1.0 makes the pipeline build video_self_attention_mask (shape (B, T_v, T_v)) internally,
58
+ # i.e. the same block-level self_attention_mask tensor inpaint/outpaint produce via a
59
+ # pixel-space conditioning_attention_mask. Group A (default) leaves it None.
60
+ GROUP_B = os.environ.get("LTX_GROUP_B", "0").strip().lower() in ("1", "true", "yes")
61
+ COND_ATTN_STRENGTH = 0.9 if GROUP_B else 1.0
62
+
63
+
64
+ def _ref_frames(n, w, h):
65
+ """Synthetic grayscale reference frames (compilation needs valid shapes, not nice pixels)."""
66
+ yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
67
+ out = []
68
+ for t in range(n):
69
+ g = (np.sin((xx / w + yy / h + t / max(n, 1)) * 2 * np.pi) * 0.5 + 0.5) * 255
70
+ out.append(Image.fromarray(g.astype(np.uint8)).convert("RGB"))
71
+ return out
72
+
73
+
74
+ def _run_pipe(steps, output_path=None):
75
+ ref = _ref_frames(NUM_FRAMES, WIDTH, HEIGHT)
76
+ out = pipe(
77
+ prompt="a colorful natural scene with gentle ambient sound",
78
+ negative_prompt="",
79
+ reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
80
+ reference_downscale_factor=1,
81
+ conditioning_attention_strength=COND_ATTN_STRENGTH,
82
+ width=WIDTH, height=HEIGHT, num_frames=NUM_FRAMES, frame_rate=FPS,
83
+ num_inference_steps=steps, sigmas=DISTILLED_SIGMA_VALUES,
84
+ guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
85
+ generator=torch.Generator(device="cuda").manual_seed(0),
86
+ output_type="np", return_dict=False,
87
+ )
88
+ if output_path is not None:
89
+ from diffusers.utils import encode_video
90
+ video_np, audio = out[0], out[1]
91
+ kw = {}
92
+ if audio is not None:
93
+ kw = dict(audio=audio[0].float().cpu(),
94
+ audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
95
+ encode_video(video_np[0], fps=FPS, output_path=output_path, **kw)
96
+ return out
97
+
98
+
99
+ def _build_dynamic_shapes(block, call):
100
+ """Flat dynamic_shapes dict. The block forward has NO **kwargs (clean diffusers
101
+ signature), so no clean-forward hack is needed. Only the video-token count T_v and
102
+ audio-token count T_a vary (text is padded to a fixed 1024 -> static). Both are large,
103
+ so size-matching is collision-safe vs structural sizes (head_dim 128, heads 32,
104
+ caption 3840, etc.). Recurse into tuples (rotary embeddings are (cos, sin) pairs)."""
105
+ import inspect
106
+ from torch.export import Dim
107
+
108
+ posnames = [n for n, p in inspect.signature(type(block).forward).parameters.items()
109
+ if n != "self" and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)]
110
+ named = {posnames[i]: a for i, a in enumerate(call.args)}
111
+ named.update(call.kwargs or {})
112
+
113
+ T_v = named["hidden_states"].shape[1]
114
+ T_a = named["audio_hidden_states"].shape[1]
115
+ DYN = {T_v, T_a}
116
+
117
+ def spec(v):
118
+ if torch.is_tensor(v):
119
+ d = {i: Dim.DYNAMIC for i, s in enumerate(v.shape) if s in DYN}
120
+ return d or None
121
+ if isinstance(v, (list, tuple)):
122
+ return type(v)(spec(x) for x in v)
123
+ return None
124
+
125
+ return {k: spec(v) for k, v in named.items()}, T_v, T_a
126
+
127
+
128
+ def compile_and_save(module: torch.nn.Module, package_dir: str):
129
+ submodule = "transformer_blocks"
130
+ block = module.get_submodule(submodule)[0]
131
+ with spaces.aoti_capture(block) as call:
132
+ _run_pipe(steps=NUM_STEPS) # aoti_capture raises at the first block call
133
+
134
+ print("AOTI: captured block forward "
135
+ f"(args={len(call.args or ())}, kwargs={sorted((call.kwargs or {}).keys())})")
136
+ for k, v in (call.kwargs or {}).items():
137
+ if torch.is_tensor(v):
138
+ print(f" {k}: Tensor {tuple(v.shape)} {v.dtype}")
139
+ elif isinstance(v, (list, tuple)):
140
+ inner = [tuple(x.shape) if torch.is_tensor(x) else type(x).__name__ for x in v]
141
+ print(f" {k}: {type(v).__name__} {inner}")
142
+ else:
143
+ print(f" {k}: {type(v).__name__} {v}")
144
+
145
+ dynamic_shapes, T_v, T_a = _build_dynamic_shapes(block, call)
146
+ def _fmt(v):
147
+ if isinstance(v, dict):
148
+ return sorted(v)
149
+ if isinstance(v, (list, tuple)):
150
+ return [_fmt(x) for x in v]
151
+ return v
152
+ print(f"AOTI: T_v={T_v} T_a={T_a}; dynamic dims="
153
+ f"{ {k: _fmt(v) for k, v in dynamic_shapes.items() if v} }")
154
+
155
+ with torch.no_grad():
156
+ exported = torch.export.export(
157
+ block, args=call.args, kwargs=call.kwargs, dynamic_shapes=dynamic_shapes,
158
+ )
159
+ print("AOTI: torch.export OK")
160
+ spaces.aoti_compile_and_save(
161
+ package_dir=package_dir, exported_program=exported, submodule=submodule,
162
+ )
163
+ print("AOTI: compile_and_save OK")
164
+
165
+
166
+ def generate_samples(samples_dir: str):
167
+ if SAMPLE_MODE != "real":
168
+ import imageio.v2 as imageio
169
+ frames = [(np.random.default_rng(i).random((64, 64, 3)) * 255).astype(np.uint8) for i in range(8)]
170
+ imageio.mimsave(f"{samples_dir}/video.mp4", frames, fps=8, macro_block_size=1)
171
+ return
172
+ _run_pipe(steps=NUM_STEPS, output_path=f"{samples_dir}/video.mp4")
173
+
174
+
175
+ def main():
176
+ create_aoti_repo(
177
+ module=pipe.transformer,
178
+ module_expr="pipe.transformer",
179
+ compile_and_save=compile_and_save,
180
+ generate_samples=generate_samples,
181
+ )
182
+
183
+
184
+ # =========================
185
+ # Internal (avoid editing) — same harness as the reference AOTI job
186
+ # =========================
187
+
188
+
189
+ import inspect
190
+ import json
191
+ import random
192
+ import shutil
193
+ import sys
194
+ import time
195
+ from packaging.version import Version
196
+ from pathlib import Path
197
+ from tempfile import TemporaryDirectory
198
+ from typing import Callable
199
+
200
+ import huggingface_hub as hf
201
+ from requests.exceptions import HTTPError
202
+
203
+
204
+ def create_aoti_repo(module, module_expr, compile_and_save, generate_samples, aoti_loader=None):
205
+ HUB_URL = 'https://huggingface.co'
206
+ user = hf.whoami()['name']
207
+ job_id = os.environ.get('JOB_ID')
208
+ job_info = hf.inspect_job(job_id=job_id) if job_id is not None else None
209
+ env_info = torch.utils.collect_env.get_env_info()
210
+ library_name, config = _get_library_config(module)
211
+
212
+ with TemporaryDirectory() as tempdir:
213
+ tempdir = Path(tempdir)
214
+ readme_path = tempdir / 'README.md'
215
+ package_dir = tempdir / 'package'
216
+ samples_before_dir = tempdir / 'samples' / 'before'
217
+ samples_after_dir = tempdir / 'samples' / 'after'
218
+ environment_path = tempdir / 'environment.json'
219
+ config_path = tempdir / 'module_config.json'
220
+
221
+ samples_before_dir.mkdir(parents=True)
222
+ t0 = time.perf_counter()
223
+ generate_samples(str(samples_before_dir))
224
+ generate_before_dt = time.perf_counter() - t0
225
+
226
+ package_dir.mkdir(parents=True)
227
+ compile_and_save(module, str(package_dir))
228
+ if aoti_loader is not None:
229
+ aoti_loader(module, str(package_dir))
230
+ else:
231
+ spaces.aoti_load_from_package_dir(module, package_dir)
232
+
233
+ samples_after_dir.mkdir(parents=True)
234
+ t0 = time.perf_counter()
235
+ generate_samples(str(samples_after_dir))
236
+ generate_after_dt = time.perf_counter() - t0
237
+
238
+ environment_path.write_text(json.dumps(env_info._asdict(), indent=4))
239
+ if config is not None:
240
+ config_path.write_text(json.dumps(config, indent=4))
241
+
242
+ output_repo_id = _create_empty_repo(
243
+ user=user, module=module, cuda_version=env_info.cuda_runtime_version,
244
+ kernels=(package_dir / 'kernels').is_dir(),
245
+ )
246
+
247
+ model_init_region = (inspect.getsource(sys.modules['__main__'])
248
+ .split('\n# README::MODEL_INIT::START')[1]
249
+ .split('\n# README::MODEL_INIT::END')[0])
250
+ aoti_load_readme = spaces.aoti_load_call_source(
251
+ module_expr=module_expr, repo_id=output_repo_id, aoti_loader=aoti_loader)
252
+ def get_link(path: Path):
253
+ kind = 'tree' if path.is_dir() else 'resolve'
254
+ return f'{HUB_URL}/{output_repo_id}/{kind}/main/{path.relative_to(tempdir)}'
255
+ readme_path.write_text(_readme_template(
256
+ model_init=model_init_region, aoti_load=aoti_load_readme, repo_id=output_repo_id,
257
+ job_id=f'{user}/{job_id}',
258
+ job_image=job_info.docker_image if job_info is not None else os.getenv('JOB_IMAGE'),
259
+ job_flavor=job_info.flavor if job_info is not None else os.getenv('JOB_FLAVOR'),
260
+ environment=torch.utils.collect_env.pretty_str(env_info),
261
+ library_name=library_name,
262
+ generate_before_dt=generate_before_dt, generate_after_dt=generate_after_dt,
263
+ samples_before_urls=[get_link(p) for p in samples_before_dir.iterdir()],
264
+ samples_after_urls=[get_link(p) for p in samples_after_dir.iterdir()],
265
+ ))
266
+ shutil.copyfile(__file__, tempdir / 'job.py')
267
+ hf.upload_folder(repo_id=output_repo_id, folder_path=tempdir)
268
+ print(f"AoT repository successfully created at: {HUB_URL}/{output_repo_id}")
269
+
270
+
271
+ def _create_empty_repo(user, module, cuda_version, kernels, max_attempts=10):
272
+ for _ in range(max_attempts):
273
+ output_repo_id = _get_repo_id(user, module, cuda_version, kernels)
274
+ try:
275
+ hf.create_repo(output_repo_id, private=True)
276
+ except HTTPError as err:
277
+ if err.response.status_code != 409:
278
+ raise
279
+ else:
280
+ return output_repo_id
281
+ raise AssertionError
282
+
283
+
284
+ def _get_repo_id(user, module, cuda_version, kernels):
285
+ if (repo_id := os.getenv('OUTPUT_REPO_ID')) is not None:
286
+ return repo_id
287
+ namespace = os.getenv('OUTPUT_REPO_NAMESPACE', user)
288
+ base_name = os.getenv('OUTPUT_REPO_BASE_NAME', module.__class__.__name__)
289
+ sm = ''.join(map(str, torch.cuda.get_device_capability()))
290
+ cu = ''.join(cuda_version.split('.')[:2])
291
+ rnd = random.randbytes(1).hex()
292
+ res = f'{namespace}/{base_name}-sm{sm}-cu{cu}'
293
+ if kernels:
294
+ torch_version = Version(torch.__version__)
295
+ res += f'-torch{torch_version.major}{torch_version.minor}'
296
+ return f'{res}-r{rnd}'
297
+
298
+
299
+ def _get_library_config(module):
300
+ if (config := getattr(module, 'config', None)) is None:
301
+ return None, None
302
+ if callable(getattr(config, 'to_dict', None)):
303
+ config = config.to_dict()
304
+ if not isinstance(config, dict):
305
+ return None, None
306
+ if 'transformers_version' in config:
307
+ library_name = 'transformers'
308
+ elif '_diffusers_version' in config:
309
+ library_name = 'diffusers'
310
+ else:
311
+ library_name = 'unknown'
312
+ return library_name, config
313
+
314
+
315
+ def _readme_template(model_init, aoti_load, repo_id, job_id, job_image, job_flavor,
316
+ environment, library_name, generate_before_dt, generate_after_dt,
317
+ samples_before_urls, samples_after_urls):
318
+ NEWLINE = '\n'
319
+ IMAGE_EXTS = ('.png', '.webp', '.jpg', '.jpeg', '.gif')
320
+ VIDEO_EXTS = ('.mp4', '.webm', '.mov')
321
+ def media_cell(url):
322
+ name = url.split('/')[-1]
323
+ if name.endswith(IMAGE_EXTS):
324
+ return f'![{name}]({url})'
325
+ if name.endswith(VIDEO_EXTS):
326
+ return f'<video src="{url}" controls></video>'
327
+ return f'[{name}]({url})'
328
+ return f"""
329
+ ---
330
+ tags:
331
+ - ahead-of-time
332
+ - pytorch
333
+ library_name: {library_name or 'pytorch'}
334
+ ---
335
+
336
+ > [!NOTE]
337
+ > This **README** has been auto-generated by the **HF Job** run linked below
338
+ > and the whole repository is a reproducible artifact of this Job
339
+
340
+ # Ahead-of-time repository
341
+
342
+ AoT repos contain **pre-compiled binaries** of PyTorch models, enabling:
343
+ - fast startup times (no `torch.compile` needed)
344
+ - significant **speedup**
345
+ - **ZeroGPU** compatibility
346
+
347
+ ## How to use
348
+ ``` python
349
+ {model_init}\n
350
+ {aoti_load}
351
+ ```
352
+
353
+ ## How to reproduce or customize
354
+ ``` bash
355
+ hf jobs uv run job.py --flavor {job_flavor or '<unknown>'} --image {job_image or '<unknown>'} --secrets HF_TOKEN
356
+ ```
357
+
358
+ ## Samples
359
+ | Before compilation ({generate_before_dt:.2f}s) | After compilation ({generate_after_dt:.2f}s) |
360
+ |---|---|
361
+ {NEWLINE.join(f"| {media_cell(b)} | {media_cell(a)} |" for b, a in zip(samples_before_urls, samples_after_urls))}
362
+
363
+ Speedup: **{generate_before_dt/generate_after_dt:.2f}x**
364
+
365
+ ## Environment
366
+ <details><summary>Click to expand</summary>
367
+
368
+ ```
369
+ {environment}
370
+ ```
371
+ </details>
372
+
373
+ ## Job run
374
+ - [{job_id}](https://huggingface.co/jobs/{job_id})
375
+ """
376
+
377
+
378
+ if __name__ == '__main__':
379
+ 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:3e1ce38500e42ee37dd48e7076e697513de061c5bf36d617967859d23e7473c6
3
+ size 2127253
samples/after/video.mp4 ADDED
Binary file (21.8 kB). View file
 
samples/before/video.mp4 ADDED
Binary file (21.8 kB). View file