# /// 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
        get_link: Callable[[Path], str] = lambda path: f'{HUB_URL}/{output_repo_id}/tree/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()