| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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)) |
| |
| |
| args = tuple(call.args) |
| kwargs = dict(call.kwargs) |
| params = list(inspect.signature(block.forward).parameters) |
| bound = {} |
| for pname, val in zip(params, args): |
| 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 = {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}") |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
