Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ 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/single_transformer_blocks/package.pt2 filter=lfs diff=lfs merge=lfs -text
+package/transformer_blocks/package.pt2 filter=lfs diff=lfs merge=lfs -text
+samples/after/edited.png filter=lfs diff=lfs merge=lfs -text
+samples/before/edited.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,172 @@
+
+---
+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 spaces
+import torch
+from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
+
+pipeline = Flux2KleinPipeline.from_pretrained(
+    'black-forest-labs/FLUX.2-klein-4B',
+    torch_dtype=torch.bfloat16
+).to('cuda')
+
+spaces.aoti_load(
+    module=pipeline.transformer,
+    repo_id='cbensimon/FLUX.2-klein-4B-sm90-cu128-glibc235-r02',
+)
+```
+
+## How to reproduce or customize
+``` bash
+# Install hf CLI
+curl -LsSf https://hf.co/cli/install.sh | bash
+
+# Login
+hf auth login
+
+# Get the job file and edit (user section) if needed
+hf download cbensimon/FLUX.2-klein-4B-sm90-cu128-glibc235-r02 job.py
+
+# Run the job and change flavor or image if needed
+hf jobs uv run job.py \
+    --flavor h200 \
+    --image pytorch/pytorch:2.9.1-cuda12.8-cudnn9-devel \
+    --secrets HF_TOKEN
+```
+
+The following job [environment variables](https://hf.co/docs/hub/en/jobs-configuration#user-defined-environment-variables)
+can be used to customize the repo name generation:
+- `OUTPUT_REPO_NAMESPACE`: taken from `HF_TOKEN` otherwise
+- `OUTPUT_REPO_BASE_NAME`: defaults to `module` class name
+- `OUTPUT_REPO_ID`: fully overtakes name generation
+
+## Samples
+
+Generated as part of the compilation job: before and after compilation
+
+| Before compilation (10.19s) | After compilation (7.62s) |
+|------------------------------------------------|----------------------------------------------|
+| ![edited.png](https://huggingface.co/cbensimon/FLUX.2-klein-4B-sm90-cu128-glibc235-r02/resolve/main/samples/before/edited.png) | ![edited.png](https://huggingface.co/cbensimon/FLUX.2-klein-4B-sm90-cu128-glibc235-r02/resolve/main/samples/after/edited.png) |
+
+Speedup: **1.34x**
+(note that this might not always reflect actual performance gain)
+
+## Environment
+<details>
+<summary>Click to expand</summary>
+
+```
+PyTorch version: 2.9.1+cu128
+Is debug build: False
+CUDA used to build PyTorch: 12.8
+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.0 (default, Oct 18 2021, 02:11:22) [Clang 13.0.0 ] (64-bit runtime)
+Python platform: Linux-6.12.64-87.122.amzn2023.x86_64-x86_64-with-glibc2.35
+Is CUDA available: True
+CUDA runtime version: 12.8.93
+CUDA_MODULE_LOADING set to: 
+GPU models and configuration: GPU 0: NVIDIA H200
+Nvidia driver version: 580.126.09
+cuDNN version: Could not collect
+Is XPU available: False
+HIP runtime version: N/A
+MIOpen runtime version: N/A
+Is XNNPACK available: True
+
+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 8488C
+CPU family:                              6
+Model:                                   143
+Thread(s) per core:                      2
+Core(s) per socket:                      48
+Socket(s):                               2
+Stepping:                                8
+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 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:                                210 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-cu12           12.8.4.1         pypi_0                pypi
+[conda] nvidia-cuda-cupti-cu12       12.8.90          pypi_0                pypi
+[conda] nvidia-cuda-nvrtc-cu12       12.8.93          pypi_0                pypi
+[conda] nvidia-cuda-runtime-cu12     12.8.90          pypi_0                pypi
+[conda] nvidia-cudnn-cu12            9.10.2.21        pypi_0                pypi
+[conda] nvidia-cufft-cu12            11.3.3.83        pypi_0                pypi
+[conda] nvidia-curand-cu12           10.3.9.90        pypi_0                pypi
+[conda] nvidia-cusolver-cu12         11.7.3.90        pypi_0                pypi
+[conda] nvidia-cusparse-cu12         12.5.8.93        pypi_0                pypi
+[conda] nvidia-cusparselt-cu12       0.7.1            pypi_0                pypi
+[conda] nvidia-nccl-cu12             2.27.5           pypi_0                pypi
+[conda] nvidia-nvjitlink-cu12        12.8.93          pypi_0                pypi
+[conda] nvidia-nvtx-cu12             12.8.90          pypi_0                pypi
+[conda] optree                       0.17.0           pypi_0                pypi
+[conda] torch                        2.9.1+cu128      pypi_0                pypi
+[conda] torchaudio                   2.9.1+cu128      pypi_0                pypi
+[conda] torchelastic                 0.2.2            pypi_0                pypi
+[conda] torchvision                  0.24.1+cu128     pypi_0                pypi
+[conda] triton                       3.5.1            pypi_0                pypi
+```
+</details>
+
+## Job run
+- [cbensimon/69cfe0474adb81dd2de74621](https://huggingface.co/jobs/cbensimon/69cfe0474adb81dd2de74621)

environment.json

@@ -0,0 +1,28 @@
+{
+    "torch_version": "2.9.1+cu128",
+    "is_debug_build": "False",
+    "cuda_compiled_version": "12.8",
+    "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.0 (default, Oct 18 2021, 02:11:22) [Clang 13.0.0 ] (64-bit runtime)",
+    "python_platform": "Linux-6.12.64-87.122.amzn2023.x86_64-x86_64-with-glibc2.35",
+    "is_cuda_available": "True",
+    "cuda_runtime_version": "12.8.93",
+    "cuda_module_loading": "",
+    "nvidia_driver_version": "580.126.09",
+    "nvidia_gpu_models": "GPU 0: NVIDIA H200",
+    "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-cu12           12.8.4.1         pypi_0                pypi\nnvidia-cuda-cupti-cu12       12.8.90          pypi_0                pypi\nnvidia-cuda-nvrtc-cu12       12.8.93          pypi_0                pypi\nnvidia-cuda-runtime-cu12     12.8.90          pypi_0                pypi\nnvidia-cudnn-cu12            9.10.2.21        pypi_0                pypi\nnvidia-cufft-cu12            11.3.3.83        pypi_0                pypi\nnvidia-curand-cu12           10.3.9.90        pypi_0                pypi\nnvidia-cusolver-cu12         11.7.3.90        pypi_0                pypi\nnvidia-cusparse-cu12         12.5.8.93        pypi_0                pypi\nnvidia-cusparselt-cu12       0.7.1            pypi_0                pypi\nnvidia-nccl-cu12             2.27.5           pypi_0                pypi\nnvidia-nvjitlink-cu12        12.8.93          pypi_0                pypi\nnvidia-nvtx-cu12             12.8.90          pypi_0                pypi\noptree                       0.17.0           pypi_0                pypi\ntorch                        2.9.1+cu128      pypi_0                pypi\ntorchaudio                   2.9.1+cu128      pypi_0                pypi\ntorchelastic                 0.2.2            pypi_0                pypi\ntorchvision                  0.24.1+cu128     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 8488C\nCPU family:                              6\nModel:                                   143\nThread(s) per core:                      2\nCore(s) per socket:                      48\nSocket(s):                               2\nStepping:                                8\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 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:                                210 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,419 @@
+# /// script
+# requires-python = "==3.10"
+# dependencies = [
+#   "spaces==0.49b5",
+#   "torch==2.9.1",
+#   "torchvision",
+#   "diffusers==0.37.1",
+#   "huggingface_hub==1.8.0",
+#   "transformers",
+#   "accelerate",
+#   "setuptools",
+# ]
+# ///
+# fmt: off
+
+
+# =========================
+# User section 
+# =========================
+
+# README::MODEL_INIT::START
+import spaces
+import torch
+from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
+
+pipeline = Flux2KleinPipeline.from_pretrained(
+    'black-forest-labs/FLUX.2-klein-4B',
+    torch_dtype=torch.bfloat16
+).to('cuda')
+# README::MODEL_INIT::END
+
+
+TRANSFORMER_IMAGE_DIM = torch.export.Dim('image_seq_length', min=4096, max=16384) # min: 0 images, max: 3 (1024x1024) images
+
+TRANSFORMER_DYNAMIC_SHAPES = {
+    'transformer_blocks': {
+        'hidden_states': {
+            1: TRANSFORMER_IMAGE_DIM,
+        },
+        'image_rotary_emb': (
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+        ),
+    },
+    'single_transformer_blocks': {
+        'hidden_states': {
+            1: TRANSFORMER_IMAGE_DIM + 512,
+        },
+        'image_rotary_emb': (
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+        ),
+    },
+}
+
+INDUCTOR_CONFIGS = {
+    'conv_1x1_as_mm': True,
+    'epilogue_fusion': False,
+    'coordinate_descent_tuning': True,
+    'coordinate_descent_check_all_directions': True,
+    'max_autotune': True,
+    'triton.cudagraphs': True,
+}
+
+
+def compile_and_save(module: torch.nn.Module, package_dir: str):
+    from PIL import Image
+    from torch.utils._pytree import tree_map
+
+    for submodule in (
+        'transformer_blocks',
+        'single_transformer_blocks',
+    ):
+        block = module.get_submodule(submodule)[0]
+        with spaces.aoti_capture(block) as call:
+            pipeline(
+                prompt="prompt",
+                image=[Image.new("RGB", (1024, 1024))],
+            )
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[submodule]
+        with torch.no_grad():
+            exported = torch.export.export(
+                mod=block,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
+            )
+        spaces.aoti_compile_and_save(
+            package_dir=package_dir,
+            exported_program=exported,
+            inductor_configs=INDUCTOR_CONFIGS,
+            submodule=submodule,
+        )
+
+
+def generate_samples(samples_dir: str):
+    from diffusers.utils.loading_utils import load_image
+    output = pipeline(
+        prompt="Remove the sunglasses",
+        image=load_image('https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/wan-cat.jpg'),
+        guidance_scale=2.5,
+        generator=torch.Generator(device='cuda').manual_seed(42),
+    )
+    output.images[0].save(f'{samples_dir}/edited.png')
+
+
+def main():
+    create_aoti_repo(
+        module=pipeline.transformer,
+        module_expr='pipeline.transformer',
+        compile_and_save=compile_and_save,
+        generate_samples=generate_samples,
+    )
+
+
+# =========================
+# Internal (avoid editing)
+# =========================
+
+
+import inspect
+import json
+import os
+import platform
+import random
+import shutil
+import sys
+import time
+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: torch.nn.Module,
+    module_expr: str,
+    compile_and_save: Callable[[torch.nn.Module, str], None],
+    generate_samples: Callable[[str], None],
+    aoti_loader: Callable[[torch.nn.Module, str], None] | None = None,
+):
+    """
+    Compile a PyTorch module Ahead-of-Time and publish it to a Hugging Face Hub repository.
+
+    This function orchestrates the following process by using the passed handlers:
+    - generates output samples
+    - compile the module and save it
+    - load the compiled version (mutates the module)
+    - generates output samples again (meaning after compilation)
+    - creates a README + other context files
+    - uploads everything to the Hub
+
+    Parameters
+    ----------
+    module:
+        Module to compile (e.g. `pipeline.transformer`).
+
+    module_expr:
+        String representation used in generated README code.
+
+    compile_and_save:
+        Callable with signature `(module: torch.nn.Module, package_dir: str) -> None`.  
+        Must AoT compile `module` (without mutating) to `package_dir` using `spaces.aoti_compile_and_save`
+
+    generate_samples:
+        Callable with signature `(samples_dir: str) -> None`.  
+        Must generate samples from the model (e.g. `pipeline`) and save them inside `samples_dir`
+
+    aoti_loader:
+        Callable with signature `(module: torch.nn.Module, package_dir: str) -> None`.  
+        Optional custom loader passed to `spaces.aoti_load`.  
+        Defaults to `spaces.aoti_load_from_package_dir`
+    """
+
+    HUB_URL = 'https://huggingface.co'
+
+    user = hf.whoami()['name']
+    job_id = os.environ['JOB_ID']
+    job_info = hf.inspect_job(job_id=job_id)
+    env_info = torch.utils.collect_env.get_env_info()
+    library_name, config = _get_library_config(module)
+    output_repo_id = _create_empty_repo(user, module, env_info.cuda_runtime_version)
+
+    print(f"Created empty repo: {HUB_URL}/{output_repo_id}")
+
+    # README: how-to-use
+    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,
+    )
+
+    with TemporaryDirectory() as tempdir:
+        tempdir = Path(tempdir)
+
+        # Structure
+        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 compile
+        samples_before_dir.mkdir(parents=True)
+        t0 = time.perf_counter()
+        generate_samples(str(samples_before_dir))
+        generate_before_dt = time.perf_counter() - t0
+
+        # Compile and load
+        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 compile
+        samples_after_dir.mkdir(parents=True)
+        t0 = time.perf_counter()
+        generate_samples(str(samples_after_dir))
+        generate_after_dt = time.perf_counter() - t0
+
+        # Environment  and config dump
+        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))
+
+        # README.md
+        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,
+            job_flavor=job_info.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(path) for path in samples_before_dir.iterdir()],
+            samples_after_urls=[get_link(path) for path in samples_after_dir.iterdir()],
+        ))
+
+        # Self include
+        shutil.copyfile(__file__, tempdir / 'job.py')
+
+        # Push to hub
+        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: str,
+    module: torch.nn.Module,
+    cuda_runtime_version: str,
+    max_attempts: int = 10
+):
+    for _ in range(max_attempts):
+        output_repo_id = _get_repo_id(user, module, cuda_runtime_version)
+        try:
+            hf.create_repo(output_repo_id)
+        except HTTPError as err:
+            if err.response.status_code != 409:
+                raise
+        else:
+            return output_repo_id
+    raise AssertionError
+
+
+def _get_repo_id(
+    user: str,
+    module: torch.nn.Module,
+    cuda_runtime_version: str,
+):
+    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_runtime_version.split('.')[:2])
+    glibc = platform.libc_ver()[1].replace('.', '')
+    rnd = random.randbytes(1).hex()
+    return f'{namespace}/{base_name}-sm{sm}-cu{cu}-glibc{glibc}-r{rnd}'
+
+
+def _get_library_config(module: torch.nn.Module):
+    if (config := getattr(module, 'config', None)) is None:
+        return None, None
+    if callable(getattr(config, 'to_dict', None)):
+        config_dict = config.to_dict()
+    elif isinstance(getattr(config, '__dict__', None), dict):
+        config_dict = config.__dict__
+    else:
+        return None, None
+    if 'transformers_version' in config_dict:
+        library_name = 'transformers'
+    elif '_diffusers_version' in config_dict:
+        library_name = 'diffusers'
+    else:
+        library_name = 'unknown'
+    return library_name, config_dict
+
+
+def _readme_template(
+    model_init: str,
+    aoti_load: str,
+    repo_id: str,
+    job_id: str,
+    job_image: str | None,
+    job_flavor: str | None,
+    environment: str,
+    library_name: str | None,
+    generate_before_dt: float,
+    generate_after_dt: float,
+    samples_before_urls: list[str],
+    samples_after_urls: list[str],
+):
+    NEWLINE = '\n'
+    IMAGE_EXTS = ('.png', '.webp', '.jpg', '.jpeg', '.gif')
+    VIDEO_EXTS = ('.mp4', '.webm', '.mov')
+
+    def media_cell(url: str):
+        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
+# Install hf CLI
+curl -LsSf https://hf.co/cli/install.sh | bash
+
+# Login
+hf auth login
+
+# Get the job file and edit (user section) if needed
+hf download {repo_id} job.py
+
+# Run the job and change flavor or image if needed
+hf jobs uv run job.py \\
+    --flavor {job_flavor or '<unknown>'} \\
+    --image {job_image or '<unknown>'} \\
+    --secrets HF_TOKEN
+```
+
+The following job [environment variables](https://hf.co/docs/hub/en/jobs-configuration#user-defined-environment-variables)
+can be used to customize the repo name generation:
+- `OUTPUT_REPO_NAMESPACE`: taken from `HF_TOKEN` otherwise
+- `OUTPUT_REPO_BASE_NAME`: defaults to `module` class name
+- `OUTPUT_REPO_ID`: fully overtakes name generation
+
+## Samples
+
+Generated as part of the compilation job: before and after compilation
+
+| Before compilation ({generate_before_dt:.2f}s) | After compilation ({generate_after_dt:.2f}s) |
+|------------------------------------------------|----------------------------------------------|
+{NEWLINE.join(
+    f"| {media_cell(before_url)} | {media_cell(after_url)} |"
+    for before_url, after_url in zip(samples_before_urls, samples_after_urls)
+)}
+
+Speedup: **{generate_before_dt/generate_after_dt:.2f}x**
+(note that this might not always reflect actual performance gain)
+
+## 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,25 @@
+{
+    "patch_size": 1,
+    "in_channels": 128,
+    "out_channels": null,
+    "num_layers": 5,
+    "num_single_layers": 20,
+    "attention_head_dim": 128,
+    "num_attention_heads": 24,
+    "joint_attention_dim": 7680,
+    "timestep_guidance_channels": 256,
+    "mlp_ratio": 3.0,
+    "axes_dims_rope": [
+        32,
+        32,
+        32,
+        32
+    ],
+    "rope_theta": 2000,
+    "eps": 1e-06,
+    "guidance_embeds": false,
+    "_class_name": "Flux2Transformer2DModel",
+    "_diffusers_version": "0.37.0.dev0",
+    "_name_or_path": "/root/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294/transformer",
+    "_FrozenDict__frozen": true
+}

package/single_transformer_blocks/package.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67ad15e4c7523016e5ed4569141f171f57e9778f57a7ee1c42312b9e9384c70d
+size 583455

package/transformer_blocks/package.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7506e85384600a93a0879a2d77f366c976be3621c13fe8349b452cf231327a89
+size 1062565