|
|
import functools |
|
|
import os |
|
|
import shutil |
|
|
import torch |
|
|
|
|
|
from pathlib import Path |
|
|
from urllib.request import Request, urlopen |
|
|
from typing import Optional |
|
|
|
|
|
|
|
|
def variant_cache_dir(): |
|
|
hf_hub_cache = os.environ.get("HF_HUB_CACHE") |
|
|
if hf_hub_cache is not None: |
|
|
return Path(hf_hub_cache) / "md_variants" |
|
|
|
|
|
hf_home = os.environ.get("HF_HOME") |
|
|
if hf_home is not None: |
|
|
return Path(hf_home) / "hub" / "md_variants" |
|
|
|
|
|
return Path("~/.cache/huggingface/hub").expanduser() / "md_variants" |
|
|
|
|
|
|
|
|
def cached_variant_path(variant_id: str): |
|
|
variant, *rest = variant_id.split("/", 1) |
|
|
step = rest[0] if rest else "final" |
|
|
|
|
|
cache_dir = variant_cache_dir() / variant |
|
|
os.makedirs(cache_dir, exist_ok=True) |
|
|
dest = cache_dir / f"{step}.pt" |
|
|
if dest.exists(): |
|
|
return dest |
|
|
|
|
|
md_endpoint = os.getenv("MOONDREAM_ENDPOINT", "https://api.moondream.ai") |
|
|
|
|
|
headers = {"User-Agent": "moondream-torch"} |
|
|
api_key = os.getenv("MOONDREAM_API_KEY") |
|
|
if api_key is not None: |
|
|
headers["X-Moondream-Auth"] = api_key |
|
|
|
|
|
req = Request(f"{md_endpoint}/v1/variants/{variant_id}/download", headers=headers) |
|
|
with urlopen(req) as r, open(dest, "wb") as f: |
|
|
shutil.copyfileobj(r, f) |
|
|
return dest |
|
|
|
|
|
|
|
|
def nest(flat): |
|
|
tree = {} |
|
|
for k, v in flat.items(): |
|
|
parts = k.split(".") |
|
|
d = tree |
|
|
for p in parts[:-1]: |
|
|
d = d.setdefault(p, {}) |
|
|
d[parts[-1]] = v |
|
|
return tree |
|
|
|
|
|
|
|
|
@functools.lru_cache(maxsize=5) |
|
|
def variant_state_dict(variant_id: Optional[str] = None, device: str = "cpu"): |
|
|
if variant_id is None: |
|
|
return None |
|
|
|
|
|
state_dict = torch.load( |
|
|
cached_variant_path(variant_id), map_location=device, weights_only=True |
|
|
) |
|
|
|
|
|
|
|
|
rename_rules = [ |
|
|
("text_model.transformer.h", "text.blocks"), |
|
|
(".mixer", ".attn"), |
|
|
(".out_proj", ".proj"), |
|
|
(".Wqkv", ".qkv"), |
|
|
(".parametrizations.weight.0", ""), |
|
|
] |
|
|
new_state_dict = {} |
|
|
for key, tensor in state_dict.items(): |
|
|
new_key = key |
|
|
for old, new in rename_rules: |
|
|
if old in new_key: |
|
|
new_key = new_key.replace(old, new) |
|
|
new_state_dict[new_key] = tensor |
|
|
|
|
|
return nest(new_state_dict) |
|
|
|
