Spaces:
Running on Zero
Running on Zero
| """Checkpoint file resolution: local file / local dir / HF Hub repo → local path.""" | |
| from pathlib import Path | |
| from typing import Optional | |
| WEIGHTS_NAME = "model.safetensors" | |
| def resolve_checkpoint_file( | |
| checkpoint: str, | |
| filename: str = WEIGHTS_NAME, | |
| required: bool = True, | |
| ) -> Optional[str]: | |
| """Resolve ``filename`` inside a checkpoint reference to a local path. | |
| ``checkpoint`` may be: | |
| - a direct file path (returned as-is when it is / stands for ``filename``), | |
| - a local checkpoint directory (looks for ``filename`` inside), | |
| - an HF Hub repo id (downloads ``filename`` to the cache; private repos | |
| are authorized via the ``HF_TOKEN`` environment variable). | |
| With ``required=False`` an absent file returns None instead of raising. | |
| """ | |
| path = Path(str(checkpoint)) | |
| if path.is_file(): | |
| if path.name == filename or (required and filename == WEIGHTS_NAME): | |
| return str(path) | |
| if required: | |
| raise FileNotFoundError(f"{checkpoint!r} is a file, not {filename}") | |
| return None | |
| if path.is_dir(): | |
| local = path / filename | |
| if local.exists(): | |
| return str(local) | |
| if required: | |
| raise FileNotFoundError(f"No {filename} in {path}") | |
| return None | |
| # Not on the local filesystem → treat as an HF Hub repo id. | |
| from huggingface_hub import hf_hub_download | |
| from huggingface_hub.errors import EntryNotFoundError | |
| try: | |
| return hf_hub_download(repo_id=str(checkpoint), filename=filename) | |
| except EntryNotFoundError: | |
| if required: | |
| raise | |
| return None | |
| def resolve_safetensors(checkpoint: str) -> str: | |
| """Local path to the checkpoint's ``model.safetensors`` (downloading if needed).""" | |
| return resolve_checkpoint_file(checkpoint, WEIGHTS_NAME, required=True) | |
| def normalize_scalar_shapes(state_dict, model) -> list: | |
| """Reconcile numel-1 parameter shapes between a checkpoint and the live model. | |
| Some published checkpoints store scalar parameters (e.g. | |
| ``ref_compressor.output_scale``) as 0-dim tensors while the model declares | |
| them shape ``[1]``. The two broadcast identically, but ``load_state_dict`` | |
| hard-errors on any shape mismatch — so every load seam calls this first. | |
| Reshapes in place (both directions, guarded to numel==1 on both sides so a | |
| real shape bug still fails loudly). Returns the reshaped keys. | |
| """ | |
| fixed = [] | |
| params = dict(model.named_parameters()) | |
| for key, tensor in state_dict.items(): | |
| target = params.get(key) | |
| if target is None: | |
| continue | |
| if tensor.shape != target.shape and tensor.numel() == 1 and target.numel() == 1: | |
| state_dict[key] = tensor.reshape(target.shape) | |
| fixed.append(key) | |
| return fixed | |