Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ 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/submodules/injection/package.pt2 filter=lfs diff=lfs merge=lfs -text
+package/submodules/plain/package.pt2 filter=lfs diff=lfs merge=lfs -text
+package/submodules/refiner/package.pt2 filter=lfs diff=lfs merge=lfs -text
+samples/aoti.webp filter=lfs diff=lfs merge=lfs -text
+samples/eager.webp filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,30 @@
+{
+  "base": "diffusers-internal-dev/Fibo-2.0-dev",
+  "arch": "sm120",
+  "torch": "2.11.0+cu130",
+  "cuda": "13.0",
+  "device": "NVIDIA RTX PRO 6000 Blackwell Server Edition",
+  "resolution": 1024,
+  "steps": 8,
+  "guidance": 5.0,
+  "injection_layer_ids": [
+    6,
+    14,
+    22,
+    30,
+    38
+  ],
+  "n_layers": 44,
+  "n_refiner_layers": 2,
+  "stacks": {
+    "plain": "dynamic",
+    "injection": "dynamic",
+    "refiner": "dynamic"
+  },
+  "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": {}
+}

job.py

@@ -0,0 +1,200 @@
+# /// 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)

package/submodules/injection/package.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:102cf5098dae0d2b9ed7798a35aa0a2be3a8abb89cbe9bcedf8ae880c371db38
+size 901648

package/submodules/plain/package.pt2

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

package/submodules/refiner/package.pt2

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