Spaces:
Running on Zero
Running on Zero
| """Model lifecycle: load at startup, report status, install decode hooks. | |
| `load_models()` is invoked once at startup by `app.py:status()`. It populates | |
| the `PIPE`, `MANAGER`, `LOADED_IN`, `TURBO_AVAILABLE`, `_TURBO_LoRA` globals | |
| in `state.py` (everything else reads from there). | |
| """ | |
| from __future__ import annotations | |
| import time | |
| import spaces # noqa: F401 — patches torch.cuda so the 72 GiB load can happen at startup | |
| from config import ( | |
| ATTENTION, | |
| CONDITIONER_SPACE, | |
| CHECKPOINT, | |
| MODEL_REPO, | |
| PLACEMENT, | |
| TURBO_LORA_FILE, | |
| TURBO_LORA_REPO, | |
| ) | |
| from pipeline import lower_duration_floor | |
| import state as _state | |
| from state import _log | |
| def status() -> str: | |
| if _state.LOAD_ERROR: | |
| return _state.LOAD_ERROR | |
| if _state.PIPE is None: | |
| return f"Loading `{CHECKPOINT}` transformer and `{MODEL_REPO}` VAEs. Watch the Space logs." | |
| import h3_aoti | |
| return ( | |
| f"Ready · `{CHECKPOINT}` · **bf16 transformer / fp32 VAEs** · placement `{PLACEMENT}` · attention `{ATTENTION}` · " | |
| f"{h3_aoti.status() if CHECKPOINT == 'base' else 'AoTI off (curve checkpoint)'} · " | |
| f"loaded in {_state.LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`" | |
| + (f" · turbo ready" if _state.TURBO_AVAILABLE else "") | |
| ) | |
| def load_models() -> str | None: | |
| """Load the denoising half at startup. | |
| `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and `video_processor`, | |
| so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` are never touched. | |
| Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio VAE decodes | |
| the soundtrack roughly 20 dB too quiet. | |
| """ | |
| if _state.PIPE is not None or _state.LOAD_ERROR is not None: | |
| return _state.LOAD_ERROR | |
| started = time.time() | |
| try: | |
| import torch | |
| from diffusers import ComponentsManager | |
| from h3_split_blocks import MiniMaxH3GeneratorBlocks | |
| lower_duration_floor() | |
| manager = ComponentsManager() | |
| blocks = MiniMaxH3GeneratorBlocks() | |
| _log(f"loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...") | |
| pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3") | |
| if CHECKPOINT != "base": | |
| from h3_checkpoint import load_transformer | |
| # Register first so the original transformer is never downloaded. | |
| pipe.update_components(transformer=load_transformer( | |
| MODEL_REPO, _log, merged_turbo=CHECKPOINT == "10eros-turbo-beta5" | |
| )) | |
| pipe.load_components(dtype=torch.bfloat16) | |
| pipe.transformer.set_attention_backend(ATTENTION) | |
| _log(f"components loaded in {time.time() - started:.0f}s") | |
| if TURBO_LORA_REPO and CHECKPOINT == "base": | |
| # Turbo weights (~850 MB) exceed the user custom-LoRA quota; load them | |
| # under a higher system limit. Failure must not abort the main model load. | |
| try: | |
| from safetensors.torch import load_file as _load_lora | |
| from lora.fetch import _fetch_lora_file | |
| _t0 = time.time() | |
| _log(f"loading turbo LoRA from {TURBO_LORA_REPO} ...") | |
| # System turbo file is ~851 MB; custom-LoRA cap is 512 MB by default. | |
| _lora_path = _fetch_lora_file( | |
| TURBO_LORA_REPO, TURBO_LORA_FILE, max_bytes=2 * 1024 * 1024 * 1024 | |
| ) | |
| _lora_sd = _load_lora(_lora_path) | |
| _lora_pairs: dict[str, dict[str, object]] = {} | |
| for _k, _v in _lora_sd.items(): | |
| if ".lora_A." in _k: | |
| _base = _k.split(".lora_A.")[0] | |
| _lora_pairs.setdefault(_base, {})["A"] = _v | |
| elif ".lora_B." in _k: | |
| _base = _k.split(".lora_B.")[0] | |
| _lora_pairs.setdefault(_base, {})["B"] = _v | |
| _state._TURBO_LoRA = [] | |
| for _base_name, _pair in _lora_pairs.items(): | |
| if "A" not in _pair or "B" not in _pair: | |
| continue | |
| _base_name = _base_name.removeprefix("transformer.") | |
| _parts = _base_name.split(".") | |
| _obj = pipe.transformer | |
| for _p in _parts: | |
| _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p) | |
| _A = _pair["A"].to(_obj.weight.dtype) | |
| _B = _pair["B"].to(_obj.weight.dtype) | |
| _state._TURBO_LoRA.append((_obj, _A, _B)) | |
| del _lora_sd, _lora_pairs | |
| import gc as _gc; _gc.collect() | |
| from lora.apply import _validate_lora_shapes | |
| _validate_lora_shapes(_state._TURBO_LoRA) | |
| _state.TURBO_AVAILABLE = bool(_state._TURBO_LoRA) | |
| _log(f"turbo LoRA parsed ({len(_state._TURBO_LoRA)} modules in {time.time() - _t0:.0f}s, merge pending)") | |
| if _state._TURBO_LoRA: | |
| _m, _A, _B = _state._TURBO_LoRA[0] | |
| _log(f" sample: module={type(_m).__name__}, " | |
| f"A.shape={_A.shape}, B.shape={_B.shape}, " | |
| f"A.device={_A.device}, dtype={_A.dtype}") | |
| except Exception as _turbo_e: | |
| _state.TURBO_AVAILABLE = False | |
| _state._TURBO_LoRA = None | |
| _log(f"turbo LoRA unavailable ({type(_turbo_e).__name__}: {_turbo_e}); continuing without it", "WARN") | |
| # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU | |
| # worker. Off unless `H3_AOTI=1`. | |
| import h3_aoti | |
| if CHECKPOINT == "base": | |
| h3_aoti.maybe_load(pipe.transformer) | |
| else: | |
| _log("AoTI disabled: the original package does not support curve-compressed AdaLN") | |
| if PLACEMENT == "pack": | |
| # Scoped to the transformer. `spaces` packs every startup-resident CUDA tensor into a second on-disk copy, | |
| # and packing all 77.3 GB busts the 150 GB storage quota; the 61.7 GB transformer alone fits. The ~10 GB of | |
| # fp32 VAEs move on the first GPU call instead. | |
| pipe.transformer.to("cuda") | |
| import gc as _gc | |
| _gc.collect() | |
| torch.cuda.empty_cache() | |
| if PLACEMENT == "offload": | |
| manager.enable_auto_cpu_offload(device="cuda") | |
| _arm_decode_hooks(pipe) | |
| _state.PIPE, _state.MANAGER = pipe, manager | |
| _state.LOADED_IN = time.time() - started | |
| _log(f"ready in {_state.LOADED_IN:.0f}s") | |
| from lora.fetch import _lora_cache_path | |
| if _lora_cache_path("probe", "probe.safetensors"): | |
| from state import _LORA_CACHE_DIR | |
| _log(f"LoRA bucket cache ready: {_LORA_CACHE_DIR}") | |
| except Exception as _e: | |
| _state.LOAD_ERROR = f"Loading {CHECKPOINT} with {MODEL_REPO} components failed: {type(_e).__name__}: {_e}" | |
| _log(f"load_models failed: {_state.LOAD_ERROR}", "ERROR") | |
| return _state.LOAD_ERROR | |
| def _arm_decode_hooks(pipe): | |
| """Force every VAE decode through the .to("cpu") path even when placement=offload. | |
| `ComponentsManager.enable_auto_cpu_offload` toggles device on `__call__` boundaries only; | |
| this hook patches `decode` directly so the last `state.get("videos")` survives a worker recycle. | |
| """ | |
| armed = False | |
| for module in (pipe.vae, getattr(pipe, "audio_vae", None)): | |
| if module is None or not hasattr(module, "decode"): | |
| continue | |
| orig_decode = module.decode | |
| def _patched_decode(self, *a, _orig=orig_decode, **k): | |
| _orig(*a, **k) | |
| try: | |
| self.to("cpu") | |
| except Exception: | |
| pass | |
| module.decode = _patched_decode.__get__(module, type(module)) | |
| armed = True | |
| if armed: | |
| module.decode = armed | |