maybe working

legacy_eva_clip/eva_vit.py

@@ -5,61 +5,20 @@
 from math import pi
 import os
 from pathlib import Path
-from typing import Optional
 import importlib.util
+import json
 
 import torch
 from torch import nn
 from einops import rearrange, repeat
 import logging
 
-_hub_spec = importlib.util.find_spec("huggingface_hub")
-if _hub_spec is not None:
-    from huggingface_hub import snapshot_download  # type: ignore
-else:  # pragma: no cover - optional dependency
-    snapshot_download = None
+from safetensors.torch import load_file as load_safetensors
 
-_DEFAULT_EVA_REPO = "jiuhai/eva_clip_vision_tower"
 XFORMERS_AVAILABLE = False  # populated later once xops import resolves
 _XFORMERS_WARNING_EMITTED = False
 
 
-def _resolve_vision_checkpoint_path(vision_tower_pretrained: Optional[str]) -> str:
-    """Determine where to load the EVA visual tower weights from."""
-
-    candidate_files: list[Path] = []
-    if vision_tower_pretrained:
-        supplied = Path(vision_tower_pretrained)
-        if supplied.is_file():
-            candidate_files.append(supplied)
-        elif supplied.is_dir():
-            for filename in (
-                "pytorch_model.bin",
-                "model.safetensors",
-                "visual.pth",
-                "visual.bin",
-            ):
-                candidate = supplied / filename
-                if candidate.exists():
-                    candidate_files.append(candidate)
-                    break
-
-    if candidate_files:
-        return str(candidate_files[0])
-
-    if snapshot_download is None:
-        raise FileNotFoundError(
-            "EVA vision weights not found locally and huggingface_hub is unavailable. "
-            "Provide --vision-tower-path pointing to a directory containing the checkpoint."
-        )
-
-    cache_dir = Path(snapshot_download(repo_id=_DEFAULT_EVA_REPO))
-    default_file = cache_dir / "pytorch_model.bin"
-    if not default_file.exists():
-        raise FileNotFoundError(f"Default EVA checkpoint not found at {default_file}.")
-    return str(default_file)
-
-
 def broadcat(tensors, dim=-1):
     num_tensors = len(tensors)
     shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
@@ -668,13 +627,54 @@ class EVAVisionTransformer(nn.Module):
 
 
 def load_state_dict(checkpoint_path: str, map_location: str = "cpu", model_key: str = "model|module|state_dict", is_openai: bool = False, skip_list: list = []):
+    path = Path(checkpoint_path)
+    if path.is_dir():
+        index_file = path / "model.safetensors.index.json"
+        if index_file.exists():
+            path = index_file
+        else:
+            for filename in (
+                "model.safetensors",
+                "pytorch_model.bin",
+                "visual.pth",
+                "visual.bin",
+            ):
+                candidate = path / filename
+                if candidate.exists():
+                    path = candidate
+                    break
+        if path.is_dir():
+            raise FileNotFoundError(f"No EVA checkpoint files found under {checkpoint_path}")
     if is_openai:
-        model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
+        model = torch.jit.load(str(path), map_location="cpu").eval()
         state_dict = model.state_dict()
         for key in ["input_resolution", "context_length", "vocab_size"]:
             state_dict.pop(key, None)
+    elif path.suffix == ".json" and path.name.endswith(".safetensors.index.json"):
+        if load_safetensors is None:
+            raise ImportError(
+                "safetensors is required to load EVA vision weights from sharded checkpoints. "
+                "Install `safetensors` or provide a .bin checkpoint."
+            )
+        with open(path, "r", encoding="utf-8") as f:
+            index_data = json.load(f)
+        weight_map = index_data.get("weight_map", {})
+        shard_cache: dict[str, dict[str, torch.Tensor]] = {}
+        state_dict = {}
+        for param_name, shard_name in weight_map.items():
+            shard_path = path.with_name(shard_name)
+            if shard_name not in shard_cache:
+                shard_cache[shard_name] = load_safetensors(str(shard_path), device=map_location)
+            state_dict[param_name] = shard_cache[shard_name][param_name]
+    elif path.suffix == ".safetensors":
+        if load_safetensors is None:
+            raise ImportError(
+                "safetensors is required to load EVA vision weights from safetensors checkpoints. "
+                "Install `safetensors` or provide a .bin checkpoint."
+            )
+        state_dict = load_safetensors(str(path), device=map_location)
     else:
-        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        checkpoint = torch.load(str(path), map_location=map_location)
         for mk in model_key.split("|"):
             if isinstance(checkpoint, dict) and mk in checkpoint:
                 state_dict = checkpoint[mk]
@@ -797,9 +797,14 @@ class EVAEncoderWrapper(nn.Module):
     def __init__(self, vision_tower_pretrained, config):
         super(EVAEncoderWrapper, self).__init__()
         self.config = config
-        vision_tower_path = _resolve_vision_checkpoint_path(vision_tower_pretrained)
-        self.config["vision_tower_path"] = vision_tower_path
-        self.vision_tower_path = vision_tower_path
+        if not vision_tower_pretrained:
+            raise ValueError("vision_tower_pretrained must be provided.")
+        vision_tower_path = Path(vision_tower_pretrained)
+        if not vision_tower_path.exists():
+            raise FileNotFoundError(f"EVA vision weights not found under {vision_tower_path}")
+        resolved_path = str(vision_tower_path.resolve())
+        self.config["vision_tower_path"] = resolved_path
+        self.vision_tower_path = resolved_path
         self.model = _build_vision_tower(**self.config)
 
     def forward(self, image, **kwargs):