feat: add local file and directory support for safetensors and GGUF#43
feat: add local file and directory support for safetensors and GGUF#43vrdn-23 wants to merge 8 commits into
Conversation
Allow `--model-id` to accept local paths (files or directories) in addition to Hugging Face Hub model IDs. Local paths are detected when the value contains a path separator and exists on disk. Supports all safetensors layouts (single, sharded, diffusers, sentence transformers) and GGUF files (single, multi-file, sharded) with full parity to the Hub path. KV cache estimation works when config.json is present locally, including on-demand HTTP fetch for referenced configs. New module `_local.py` provides local I/O functions that return the same data shapes as the HTTP fetch functions, so all existing parsing and display code is reused unchanged. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
@alvarobartt this is ready for a review now. I've tested this against most of scenarios and had to refactor some code into helper functions to avoid a lot of duplication. Apologies if this may be against the contributing guidelines. Please let me know if you have any thoughts/concerns/feedback |
|
@alvarobartt just wanted to ping and see if you got a chance to review this? |
There was a problem hiding this comment.
Pull request overview
Adds support for using local files/directories as --model-id, enabling memory estimation for on-disk Safetensors and GGUF models while reusing the existing parsing/printing pipeline.
Changes:
- Added local path detection in
arun()with support for local.safetensors,.gguf, and model directories (including sharded/layout variants). - Refactored shared logic into private helpers for token resolution, GGUF path filtering/result building, SentenceTransformer metadata wrapping, and KV-cache estimation.
- Updated CLI help text to document local path usage.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
| src/hf_mem/run.py | Adds local model loading path + shared helper refactors for GGUF/Safetensors handling and KV-cache estimation. |
| src/hf_mem/cli.py | Updates --model-id help to include local file/directory usage and detection rules. |
| src/hf_mem/_local.py | Introduces local I/O utilities for listing files and reading Safetensors/GGUF/JSON data. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| with open(filepath, "rb") as f: | ||
| size_bytes = f.read(8) | ||
| if len(size_bytes) < 8: | ||
| raise RuntimeError(f"File too small to be a valid safetensors file: {filepath}") | ||
| metadata_size = struct.unpack("<Q", size_bytes)[0] | ||
| metadata_bytes = f.read(metadata_size) | ||
| if len(metadata_bytes) < metadata_size: |
| # Remove __metadata__ key to match the shape returned by HTTP fetch | ||
| # (the HTTP path returns the full dict including __metadata__, and | ||
| # parse_safetensors_metadata skips it via the `if key in {"__metadata__"}` check) |
| def read_gguf_bytes(filepath: str) -> bytes: | ||
| """Read GGUF file header bytes (up to 100MB) for metadata parsing.""" | ||
| file_size = os.path.getsize(filepath) | ||
| read_size = min(file_size, _MAX_GGUF_READ_SIZE) | ||
| with open(filepath, "rb") as f: | ||
| return f.read(read_size) |
| "User-Agent": f"hf-mem/{__version__}; id={uuid4()}; model_id={model_id}; revision={revision}" | ||
| } | ||
| # --- Local path detection --- | ||
| is_local = (os.sep in model_id or model_id.startswith(".")) and os.path.exists(model_id) |
| if hf_token is not None: | ||
| return {"Authorization": f"Bearer {hf_token}"} | ||
| if token := os.getenv("HF_TOKEN"): | ||
| return {"Authorization": f"Bearer {token}"} | ||
| hf_home = os.getenv("HF_HOME", ".cache/huggingface") | ||
| token_path = ( | ||
| os.path.join(os.path.expanduser("~"), hf_home, "token") | ||
| if not os.path.isabs(hf_home) | ||
| else os.path.join(hf_home, "token") | ||
| ) |
| if prefix_match := re.match(r"(.+)-\d+-of-\d+\.gguf$", gguf_file): | ||
| prefix = prefix_match.group(1) | ||
| gguf_paths = [ | ||
| path for path in gguf_paths if re.match(rf"{re.escape(prefix)}-\d+-of-\d+\.gguf$", str(path)) |
| def list_local_files(directory: str) -> List[str]: | ||
| """Walk directory and return relative file paths, following symlinks.""" | ||
| file_paths = [] | ||
| for root, _dirs, files in os.walk(directory, followlinks=True): | ||
| for f in files: | ||
| full_path = os.path.join(root, f) | ||
| # Skip broken symlinks | ||
| if not os.path.exists(full_path): | ||
| continue | ||
| rel_path = os.path.relpath(full_path, directory) | ||
| file_paths.append(rel_path) | ||
| return file_paths |
alvarobartt
left a comment
alvarobartt
left a comment
There was a problem hiding this comment.
Hey @vrdn-23 thanks for the change, and apologies I get back to this just now, I tried to prioritize fixes (particularly around the KV cache estimations for certain quantization configurations).
I'd try to simplify it a bit so that it plays a bit nicer with the existing coding style and if possible reduce the lines-of-code it adds, as it feels a bit too much. Also some code-comments are "sloppy" and not required, so I'd remove those too e.g., the last code comment in read_safetensors_header feels unnecessary.
Thanks again, and happy to help out if required!
P.S. Could you also pull the latest changes from main? I plan to release v0.5.3 later today with those fixes, but would love to include this PR for v0.6.0 🎉
Signed-off-by: Vinay Damodaran <vrdn@hey.com>
Signed-off-by: Vinay Damodaran <vrdn@hey.com>
Signed-off-by: Vinay Damodaran <vrdn@hey.com>
|
Hey @alvarobartt , I really did try reduce the line count, but not sure if there is more I can do, without just duplicating a bunch of code all over the place to make it seem more simple. I've added you as a collaborator to my repo, so that you can take up the branches and work on the changes if there are any stylistic changes or other things you want to fix (will probably be faster than the back and forth). But if you don't have bandwidth, please feel free to leave comments and I can get to them. P.S - The comments from CoPilot are probably applicable to the remote case too for most of them. I haven't fixed any of them cause I would have to touch the remote calling part of the code too. Let me know if you would like me to address those with the existing behavior as well. |
Summary
Fixes #2
Add support for local file and directory paths in
--model-id, so users can estimate memory for models already on disk without fetching from the Hub.Local file support (
4d25c71).and exists on disk. Use./prefix for relative paths.model.safetensors.index.json), Diffusers (model_index.json), and Sentence Transformers (with Dense module detection viamodules.json).--gguf-filefiltering work identically to the remote path.config.json, including on-demand HTTP fetch for referenced configs whentext_config._name_or_pathpoints to a different Hub model (e.g. conditional generation models).src/hf_mem/_local.pywith local I/O functions (read_safetensors_header,read_local_json,read_gguf_bytes,list_local_files) that return the same data shapes as the HTTP fetches, so all existing parsing and display code is reused unchanged.--model-idhelp text updated to document local path usage.Refactor: deduplicate local/remote paths (
297aec9)The initial local support duplicated ~340 lines of logic from the remote (Hub) code path. Extracted 5 shared private helpers in
run.pyto eliminate the duplication:_resolve_hf_token: unifies 3 copies of HF token resolution (explicit param →HF_TOKENenv →$HF_HOME/tokenfile). Also removes a dead-code guard (elif "Authorization" not in headerson a freshly-created dict)._filter_gguf_paths: unifies 2 copies of--gguf-filepath filtering (shard prefix matching, exact suffix matching, error messages). Error context parameterized via alabelstring._build_gguf_result: unifies 2 copies ofResultconstruction from agguf_files_dict(single-file and multi-file sub-paths)._wrap_sentence_transformer_metadata: unifies 4 occurrences of the sentence-transformer component naming pattern (0_Transformer+ dense modules vs plainTransformer)._estimate_safetensors_kv_cache: unifies 2 copies (~60 lines each) of the KV cache estimation block. Uses anasync Callableparameter to abstract the one difference between paths — how the referenced config is fetched (local creates a one-off httpx client, remote reuses the existing one). Also improves resource safety withtry/finallyon the local path's client cleanup.Net result: ~47 fewer lines, single source of truth for all shared logic, both branches reduced to data-fetching + helper calls.