# /// 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
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}")
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''
if name.endswith(VIDEO_EXTS):
return f''
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 ''} --image {job_image or ''} --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
Click to expand
```
{environment}
```