forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
230
tinygrad_repo/extra/huggingface_onnx/huggingface_manager.py
Normal file
230
tinygrad_repo/extra/huggingface_onnx/huggingface_manager.py
Normal file
@@ -0,0 +1,230 @@
|
||||
import yaml
|
||||
import time
|
||||
import requests
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from huggingface_hub import list_models, HfApi, snapshot_download
|
||||
from tinygrad.helpers import _ensure_downloads_dir
|
||||
DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
|
||||
from tinygrad.helpers import tqdm
|
||||
|
||||
def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_dir: str|Path|None=None,
|
||||
tries: int=2, **kwargs) -> Path:
|
||||
for attempt in range(tries):
|
||||
try:
|
||||
return Path(snapshot_download(
|
||||
repo_id=repo_id,
|
||||
allow_patterns=allow_patterns,
|
||||
local_dir=str(local_dir) if local_dir is not None else None,
|
||||
**kwargs
|
||||
))
|
||||
except Exception as e:
|
||||
if attempt == tries-1: raise
|
||||
time.sleep(1)
|
||||
|
||||
# Constants for filtering models
|
||||
HUGGINGFACE_URL = "https://huggingface.co"
|
||||
SKIPPED_FILES = [
|
||||
"fp16", "int8", "uint8", "quantized", # numerical accuracy issues
|
||||
"avx2", "arm64", "avx512", "avx512_vnni", # numerical accuracy issues
|
||||
"q4", "q4f16", "bnb4", # unimplemented quantization
|
||||
"model_O4", # requires non cpu ort runner and MemcpyFromHost op
|
||||
"merged", # TODO implement attribute with graph type and Loop op
|
||||
]
|
||||
|
||||
SKIPPED_REPO_PATHS = [
|
||||
# Invalid model-index
|
||||
"AdamCodd/vit-base-nsfw-detector",
|
||||
# TODO: implement attribute with graph type and Loop op
|
||||
"minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
|
||||
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
|
||||
"HuggingFaceTB/SmolLM2-360M-Instruct",
|
||||
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
|
||||
"HuggingFaceTB/SmolLM2-1.7B-Instruct",
|
||||
# TODO: implement RandomNormalLike
|
||||
"stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
|
||||
# TODO: implement NonZero
|
||||
"mangoapps/fb_zeroshot_mnli_onnx",
|
||||
# TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
|
||||
"briaai/RMBG-2.0",
|
||||
]
|
||||
|
||||
|
||||
class HuggingFaceONNXManager:
|
||||
def __init__(self):
|
||||
self.base_dir = Path(__file__).parent
|
||||
self.models_dir = DOWNLOADS_DIR
|
||||
self.api = HfApi()
|
||||
|
||||
def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
|
||||
print(f"Discovering top {limit} ONNX models sorted by {sort}...")
|
||||
repos = []
|
||||
i = 0
|
||||
|
||||
for model in list_models(filter="onnx", sort=sort):
|
||||
if model.id in SKIPPED_REPO_PATHS:
|
||||
continue
|
||||
|
||||
print(f" {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
|
||||
repos.append(model.id)
|
||||
i += 1
|
||||
if i == limit:
|
||||
break
|
||||
|
||||
print(f"Found {len(repos)} suitable ONNX models")
|
||||
return repos
|
||||
|
||||
def collect_metadata(self, repos: list[str]) -> dict:
|
||||
print(f"Collecting metadata for {len(repos)} repositories...")
|
||||
metadata = {"repositories": {}}
|
||||
total_size = 0
|
||||
|
||||
for repo in tqdm(repos, desc="Collecting metadata"):
|
||||
try:
|
||||
files_metadata = []
|
||||
model_info = self.api.model_info(repo)
|
||||
|
||||
for file in model_info.siblings:
|
||||
filename = file.rfilename
|
||||
if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
|
||||
continue
|
||||
if any(skip_str in filename for skip_str in SKIPPED_FILES):
|
||||
continue
|
||||
|
||||
# Get file size from API or HEAD request
|
||||
try:
|
||||
head = requests.head(
|
||||
f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
|
||||
allow_redirects=True,
|
||||
timeout=10
|
||||
)
|
||||
file_size = file.size or int(head.headers.get('Content-Length', 0))
|
||||
except requests.RequestException:
|
||||
file_size = file.size or 0
|
||||
|
||||
files_metadata.append({
|
||||
"file": filename,
|
||||
"size": f"{file_size/1e6:.2f}MB"
|
||||
})
|
||||
total_size += file_size
|
||||
|
||||
if files_metadata: # Only add repos with valid ONNX files
|
||||
metadata["repositories"][repo] = {
|
||||
"url": f"{HUGGINGFACE_URL}/{repo}",
|
||||
"download_path": None,
|
||||
"files": files_metadata,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
print(f"WARNING: Failed to collect metadata for {repo}: {e}")
|
||||
continue
|
||||
|
||||
metadata['total_size'] = f"{total_size/1e9:.2f}GB"
|
||||
metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
|
||||
print(f"Collected metadata for {len(metadata['repositories'])} repositories")
|
||||
print(f"Total estimated download size: {metadata['total_size']}")
|
||||
|
||||
return metadata
|
||||
|
||||
def download_models(self, metadata: dict) -> dict:
|
||||
self.models_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
repos = metadata["repositories"]
|
||||
n = len(repos)
|
||||
|
||||
print(f"Downloading {n} repositories to {self.models_dir}...")
|
||||
|
||||
for i, (model_id, model_data) in enumerate(repos.items()):
|
||||
print(f" Downloading {i+1}/{n}: {model_id}...")
|
||||
|
||||
try:
|
||||
# Download ONNX model files
|
||||
allow_patterns = [file_info["file"] for file_info in model_data["files"]]
|
||||
root_path = snapshot_download_with_retry(
|
||||
repo_id=model_id,
|
||||
allow_patterns=allow_patterns,
|
||||
local_dir=str(self.models_dir / model_id)
|
||||
)
|
||||
|
||||
# Download config files (usually small)
|
||||
snapshot_download_with_retry(
|
||||
repo_id=model_id,
|
||||
allow_patterns=["*config.json"],
|
||||
local_dir=str(self.models_dir / model_id)
|
||||
)
|
||||
|
||||
model_data["download_path"] = str(root_path)
|
||||
print(f" Downloaded to: {root_path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ERROR: Failed to download {model_id}: {e}")
|
||||
model_data["download_path"] = None
|
||||
continue
|
||||
|
||||
successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
|
||||
print(f"Successfully downloaded {successful_downloads}/{n} repositories")
|
||||
print(f"All models saved to: {self.models_dir}")
|
||||
|
||||
return metadata
|
||||
|
||||
def save_metadata(self, metadata: dict, output_file: str):
|
||||
yaml_path = self.base_dir / output_file
|
||||
with open(yaml_path, 'w') as f:
|
||||
yaml.dump(metadata, f, sort_keys=False)
|
||||
print(f"Metadata saved to: {yaml_path}")
|
||||
|
||||
def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
|
||||
sort: str = "downloads", download: bool = True):
|
||||
print(f"Starting HuggingFace ONNX workflow...")
|
||||
print(f" Limit: {limit} models")
|
||||
print(f" Sort by: {sort}")
|
||||
print(f" Download: {'Yes' if download else 'No'}")
|
||||
print(f" Output: {output_file}")
|
||||
print("-" * 50)
|
||||
|
||||
repos = self.discover_models(limit, sort)
|
||||
|
||||
metadata = self.collect_metadata(repos)
|
||||
|
||||
if download:
|
||||
metadata = self.download_models(metadata)
|
||||
|
||||
self.save_metadata(metadata, output_file)
|
||||
|
||||
print("-" * 50)
|
||||
print("Workflow completed successfully!")
|
||||
if download:
|
||||
successful = sum(1 for repo in metadata["repositories"].values()
|
||||
if repo["download_path"] is not None)
|
||||
print(f"{successful}/{len(metadata['repositories'])} models downloaded")
|
||||
|
||||
return metadata
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
|
||||
)
|
||||
|
||||
parser.add_argument("--limit", type=int, help="Number of top repositories to process")
|
||||
parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
|
||||
help="Output YAML file name (default: huggingface_repos.yaml)")
|
||||
parser.add_argument("--sort", type=str, default="downloads",
|
||||
choices=["downloads", "likes", "created", "modified"],
|
||||
help="Sort criteria for model discovery (default: downloads)")
|
||||
|
||||
parser.add_argument("--download", action="store_true", default=False,
|
||||
help="Download models after collecting metadata")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.limit: parser.error("--limit is required")
|
||||
|
||||
manager = HuggingFaceONNXManager()
|
||||
manager.discover_and_download(
|
||||
limit=args.limit,
|
||||
output_file=args.output,
|
||||
sort=args.sort,
|
||||
download=args.download
|
||||
)
|
||||
Reference in New Issue
Block a user