job.py

# /// 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"


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,
        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'![{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()