Spaces:
Running on Zero
Running on Zero
| """Load aokit AOT-compiled block graphs onto a BriaFibo2Transformer2DModel, on ZeroGPU. | |
| The transformer's `layers` ModuleList is heterogeneous: injection layers (ids in the model config) | |
| carry gated cross-attention to the text encoder, the rest do not. So one graph cannot cover the whole | |
| ModuleList (as in the Flux/klein case). Instead the package holds three weight-less graphs -- | |
| `plain`, `injection`, `refiner` -- and we patch each block with the matching one. Each patched block | |
| keeps its own eager `forward` as a fallback, so a shape/version the graph was not built for still runs. | |
| Package (built by job_fibo2_aoti.py, see the AOTI repo): config.json + package/submodules/<name>/package.pt2. | |
| The compiled graphs are weight-less, so the SAME package serves dev and turbo (identical architecture). | |
| """ | |
| import json | |
| import os | |
| from pathlib import Path | |
| import torch | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| AOTI_REPO = os.environ.get("FIBO2_AOTI_REPO", "diffusers-internal-dev/Fibo-2.0-dev-aoti-sm120") | |
| AOTI_ENABLED = os.environ.get("AOTI", "1") != "0" | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| def _with_eager_fallback(compiled, eager, label): | |
| warned = [] | |
| def forward(*args, **kwargs): | |
| try: | |
| return compiled(*args, **kwargs) | |
| except Exception as e: # shape/dtype the graph was not built for -> run eager | |
| if not warned: | |
| warned.append(True) | |
| print(f"AOTI {label} fell back to eager: {e!r}", flush=True) | |
| return eager(*args, **kwargs) | |
| return forward | |
| def load_aoti(transformer): | |
| """Patch `transformer` in place with the compiled graphs. Raises if the package is | |
| unavailable or built for a different torch; the caller should fall back to eager.""" | |
| from aokit import LazyAOTIModel, patch | |
| cfg = json.loads(Path(hf_hub_download(AOTI_REPO, "config.json", token=HF_TOKEN)).read_text()) | |
| if cfg.get("torch") != torch.__version__: | |
| raise RuntimeError(f"AOTI built for torch {cfg.get('torch')}, runtime has {torch.__version__}") | |
| pkg = Path(snapshot_download(AOTI_REPO, allow_patterns="package/*", token=HF_TOKEN)) / "package" | |
| inj = set(cfg["injection_layer_ids"]) | |
| models = { | |
| name: LazyAOTIModel(pkg / "submodules" / name / "package.pt2") | |
| for name in cfg["stacks"] | |
| if (pkg / "submodules" / name / "package.pt2").exists() | |
| } | |
| def bind(block, model, label): | |
| eager = block.forward | |
| patch(block, model) # sets block.forward to the compiled graph (bound to block's own weights) | |
| block.forward = _with_eager_fallback(block.forward, eager, label) | |
| def bind_refiner(block, m_t2i, m_edit, label): | |
| # The refiner is called with a different arg set on t2i vs edit; try the t2i graph, | |
| # then the edit graph, then eager. Each graph is bound to this block's own weights. | |
| eager = block.forward | |
| patch(block, m_t2i) | |
| f_t2i = block.forward | |
| f_edit = m_edit.with_weights(block.state_dict()) | |
| warned = [] | |
| def forward(*args, **kwargs): | |
| try: | |
| return f_t2i(*args, **kwargs) | |
| except Exception: | |
| try: | |
| return f_edit(*args, **kwargs) | |
| except Exception as e: | |
| if not warned: | |
| warned.append(True) | |
| print(f"AOTI {label} fell back to eager: {e!r}", flush=True) | |
| return eager(*args, **kwargs) | |
| block.forward = forward | |
| if "plain" in models or "injection" in models: | |
| for i, block in enumerate(transformer.layers): | |
| model = models.get("injection") if i in inj else models.get("plain") | |
| if model is not None: | |
| bind(block, model, f"layers.{i}") | |
| if "refiner" in models: | |
| redit = models.get("refiner_edit") | |
| for block in transformer.noise_refiner: | |
| if redit is not None: | |
| bind_refiner(block, models["refiner"], redit, "noise_refiner") | |
| else: | |
| bind(block, models["refiner"], "noise_refiner") | |
| print(f"AOTI loaded: {list(models)} from {AOTI_REPO}", flush=True) | |