Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| from typing import Any, Dict | |
| import torch | |
| from horizonstream.models.horizonstream import HorizonStream | |
| from horizonstream.utils.hub import resolve_checkpoint_path | |
| class HorizonStreamModel(torch.nn.Module): | |
| def __init__(self, cfg: Dict[str, Any] | None): | |
| super().__init__() | |
| cfg = cfg or {} | |
| ckpt_path = resolve_checkpoint_path( | |
| cfg.get("checkpoint", None), | |
| cfg.get("hf", None), | |
| ) | |
| model_cfg = dict(cfg.get("horizonstream_cfg", {}) or {}) | |
| self.horizonstream = HorizonStream(**model_cfg) | |
| if ckpt_path: | |
| self.load_checkpoint(ckpt_path, strict=bool(cfg.get("strict_load", True))) | |
| def load_checkpoint(self, ckpt_path: str, strict: bool = True): | |
| if not os.path.exists(ckpt_path): | |
| raise FileNotFoundError(ckpt_path) | |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) | |
| if isinstance(ckpt, dict): | |
| if "model" in ckpt and isinstance(ckpt["model"], dict): | |
| state = ckpt["model"] | |
| elif "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict): | |
| state = ckpt["state_dict"] | |
| else: | |
| state = ckpt | |
| else: | |
| raise TypeError("Unsupported checkpoint format") | |
| def strip_wrappers(key: str) -> str: | |
| # Training checkpoints may be wrapped by several container prefixes. | |
| # Strip them repeatedly so we can match both wrapper and bare module states. | |
| prefixes = ("module.", "model.", "state_dict.", "sampler.") | |
| changed = True | |
| while changed: | |
| changed = False | |
| for prefix in prefixes: | |
| if key.startswith(prefix): | |
| key = key[len(prefix) :] | |
| changed = True | |
| return key | |
| normalized_state = {strip_wrappers(k): v for k, v in state.items()} | |
| full_keys = set(self.state_dict().keys()) | |
| core_keys = set(self.horizonstream.state_dict().keys()) | |
| full_matches = sum(1 for k in normalized_state.keys() if k in full_keys) | |
| core_matches = sum(1 for k in normalized_state.keys() if k in core_keys) | |
| if full_matches == 0 and core_matches > 0: | |
| # Checkpoint is likely dumped from sampler/core module directly. | |
| # Remove an optional horizonstream. prefix if it still exists. | |
| core_state = { | |
| k.removeprefix("horizonstream."): v for k, v in normalized_state.items() | |
| } | |
| missing, unexpected = self.horizonstream.load_state_dict(core_state, strict=False) | |
| else: | |
| missing, unexpected = self.load_state_dict(normalized_state, strict=False) | |
| if missing or unexpected: | |
| msg = ( | |
| "checkpoint mismatch: " | |
| f"missing={len(missing)} unexpected={len(unexpected)} " | |
| f"(full_matches={full_matches}, core_matches={core_matches})" | |
| ) | |
| if missing: | |
| sample_missing = ", ".join(missing[:20]) | |
| msg += f"\n missing(sample up to 20): {sample_missing}" | |
| if unexpected: | |
| sample_unexpected = ", ".join(unexpected[:20]) | |
| msg += f"\n unexpected(sample up to 20): {sample_unexpected}" | |
| if strict: | |
| raise RuntimeError(msg) | |
| print(msg) | |
| def forward_window(self, images: torch.Tensor): | |
| return self.horizonstream.forward_window(images) | |
| def build_sequence_state(self): | |
| return self.horizonstream.build_sequence_state() | |
| def advance_sequence_state(self, state, *, is_last_chunk: bool): | |
| self.horizonstream.advance_sequence_state(state, is_last_chunk=is_last_chunk) | |
| def forward_chunk( | |
| self, | |
| images: torch.Tensor, | |
| *, | |
| window_size: int, | |
| chunk_idx: int, | |
| state, | |
| ): | |
| return self.horizonstream.forward_chunk( | |
| images, | |
| window_size=window_size, | |
| chunk_idx=chunk_idx, | |
| frame_kv_caches=state["frame_kv_caches"], | |
| global_kv_caches=state["global_kv_caches"], | |
| gla_cache=state.get("gla_cache"), | |
| ) | |
| def forward(self, images: torch.Tensor): | |
| return self.forward_window(images) | |