# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "aokit",
#   "torch==2.11.0",
#   "torchvision==0.26.0",
#   "diffusers @ https://github.com/muhammadali-afridi/diffusers/archive/60f95274bf07ed705d675eb09d99b6ae125b601e.tar.gz",
#   "transformers==5.20.1",
#   "accelerate",
#   "safetensors",
#   "sentencepiece",
#   "huggingface_hub>=1.0",
#   "Pillow",
# ]
#
# [[tool.uv.index]]
# url = "https://download.pytorch.org/whl/cu130"
# ///
#
# Ahead-of-time compile the repeated transformer blocks of BriaFibo2Transformer2DModel
# with aokit (regional compilation), verify correctness, and push a PRIVATE package repo.
#
# Launch:
#   hf jobs uv run job_fibo2_aoti.py \
#       --flavor rtx-pro-6000 \
#       --image pytorch/pytorch:2.9.1-cuda13.0-cudnn9-devel \
#       --namespace diffusers-internal-dev \
#       --secrets HF_TOKEN \
#       -e FIBO2_MODEL=diffusers-internal-dev/Fibo-2.0-dev \
#       -e OUTPUT_REPO_ID=diffusers-internal-dev/Fibo-2.0-dev-aoti-sm120 \
#       -e STEPS=8

import inspect
import json
import os
import time
import traceback
from pathlib import Path

import torch
from huggingface_hub import HfApi
from PIL import Image

import aokit
from aokit.exporting import capture

MODEL = os.environ.get("FIBO2_MODEL", "diffusers-internal-dev/Fibo-2.0-dev")
OUT_REPO = os.environ["OUTPUT_REPO_ID"]
TOKEN = os.environ.get("HF_TOKEN")
STEPS = int(os.environ.get("STEPS", "8"))
GUIDANCE = float(os.environ.get("GUIDANCE", "5.0"))
PROMPT = os.environ.get("PROMPT", "a red bicycle leaning against a white brick wall on a sunny morning.")
H = W = 1024
DTYPE = torch.bfloat16

print(f"torch {torch.__version__} | cuda {torch.version.cuda} | device {torch.cuda.get_device_name()}", flush=True)

from diffusers import BriaFibo2Pipeline

pipe = BriaFibo2Pipeline.from_pretrained(MODEL, dtype=DTYPE, token=TOKEN).to("cuda")
pipe.set_progress_bar_config(disable=True)
t = pipe.transformer
INJ = list(getattr(t.config, "injection_layer_ids", (6, 14, 22, 30, 38)))
n_layers = len(t.layers)
plain_idx = next(i for i in range(n_layers) if i not in INJ)
inj_idx = INJ[0]
print(f"n_layers={n_layers} injection={INJ} plain_rep={plain_idx} inj_rep={inj_idx}", flush=True)


def gen(path, label):
    img = pipe(PROMPT, num_inference_steps=STEPS, guidance_scale=GUIDANCE,
               height=H, width=W, generator=torch.manual_seed(0)).images[0]
    img.save(path)
    print(f"sample[{label}] -> {path}", flush=True)
    return img


def export_block(block, name):
    """Capture inputs for `block` from a real pipeline step, then export with the
    sequence axes marked dynamic (batch kept static at the captured size)."""
    with capture(block) as call:
        pipe(PROMPT, num_inference_steps=1, guidance_scale=GUIDANCE,
             height=H, width=W, generator=torch.manual_seed(0))
    # Preserve the EXACT calling convention the transformer uses (positional args stay
    # positional) so the saved graph's in_spec matches the runtime block(...) call.
    args = tuple(call.args)
    kwargs = dict(call.kwargs)
    params = list(inspect.signature(block.forward).parameters)
    bound = {}
    for pname, val in zip(params, args):  # name the positional args for spec + dynamic_shapes
        bound[pname] = val
    bound.update(kwargs)
    x = bound["x"]
    B, S = x.shape[0], x.shape[1]
    enc = bound.get("encoder_hidden_states")
    E = enc.shape[1] if torch.is_tensor(enc) else None
    print(f"  [{name}] captured B={B} S={S} E={E} "
          f"args={[type(a).__name__ for a in args]} kw={[k for k in kwargs]}", flush=True)

    seq = torch.export.Dim("seq", min=16, max=16384)
    enc_dim = torch.export.Dim("enc", min=16, max=16384) if E and E != S else None
    size_to_dim = {S: seq}
    if enc_dim is not None:
        size_to_dim[E] = enc_dim

    def spec(v):
        if not torch.is_tensor(v):
            return None
        d = {i: size_to_dim[sz] for i, sz in enumerate(v.shape) if i != 0 and sz in size_to_dim}
        return d or None

    # dynamic_shapes keyed by argument name (torch.export maps names to args+kwargs).
    dynamic_shapes = {pname: spec(val) for pname, val in bound.items()}
    with torch.no_grad():
        ep = torch.export.export(block, args, kwargs, dynamic_shapes=dynamic_shapes)
    return ep


STACKS = {
    "plain": t.layers[plain_idx],
    "injection": t.layers[inj_idx],
    "refiner": t.noise_refiner[0],
}

api = HfApi(token=TOKEN)
work = Path("/tmp/fibo2_aoti")
pkg = work / "package"
samples = work / "samples"
pkg.mkdir(parents=True, exist_ok=True)
samples.mkdir(parents=True, exist_ok=True)

gen(samples / "eager.webp", "eager")

compiled = {}
errors = {}
for name, block in STACKS.items():
    print(f"=== export+compile stack '{name}' ===", flush=True)
    try:
        t0 = time.perf_counter()
        ep = export_block(block, name)
        aokit.compile_and_save(package_dir=str(pkg), exported_program=ep, submodule=name)
        compiled[name] = True
        print(f"  OK '{name}' in {time.perf_counter()-t0:.0f}s", flush=True)
    except Exception as e:
        errors[name] = repr(e)
        print(f"  FAILED '{name}': {e!r}", flush=True)
        traceback.print_exc()

if not compiled:
    raise RuntimeError(f"No stack compiled. Errors: {errors}")

# Load compiled stacks back with a per-index custom mapping and verify.
from aokit import LazyAOTIModel, patch

models = {name: LazyAOTIModel(pkg / "submodules" / name / "package.pt2") for name in compiled}
if "plain" in models or "injection" in models:
    for i, blk in enumerate(t.layers):
        m = models.get("injection") if i in INJ else models.get("plain")
        if m is not None:
            patch(blk, m)
if "refiner" in models:
    for blk in t.noise_refiner:
        patch(blk, models["refiner"])

gen(samples / "aoti.webp", "aoti")

cap = "".join(map(str, torch.cuda.get_device_capability()))
cu = "".join((torch.version.cuda or "0.0").split(".")[:2])
config = {
    "base": MODEL,
    "arch": f"sm{cap}",
    "torch": torch.__version__,
    "cuda": torch.version.cuda,
    "device": torch.cuda.get_device_name(),
    "resolution": H,
    "steps": STEPS,
    "guidance": GUIDANCE,
    "injection_layer_ids": INJ,
    "n_layers": n_layers,
    "n_refiner_layers": len(t.noise_refiner),
    "stacks": {name: "dynamic" for name in compiled},
    "stack_targets": {
        "plain": "layers[i] for i not in injection_layer_ids",
        "injection": "layers[i] for i in injection_layer_ids",
        "refiner": "noise_refiner[*]",
    },
    "errors": errors,
}
(work / "config.json").write_text(json.dumps(config, indent=2))
print("CONFIG:\n" + json.dumps(config, indent=2), flush=True)

# self-include this script
import shutil
main_file = getattr(__import__("__main__"), "__file__", None)
if main_file:
    shutil.copyfile(main_file, work / "job.py")

api.create_repo(OUT_REPO, repo_type="model", private=True, exist_ok=True)
api.upload_folder(repo_id=OUT_REPO, folder_path=str(work), repo_type="model")
print(f"DONE -> https://huggingface.co/{OUT_REPO}  (private)  compiled={list(compiled)} errors={list(errors)}", flush=True)