mrleo1nid/gpu-rent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Hugging Face support and enhance model resolution logic

Browse Source 
- Introduced support for Hugging Face API integration, allowing fallback model resolution when Civitai fails.
- Updated configuration to include `HF_TOKEN` and `HF_TOKEN_PATH` for authentication.
- Enhanced model capture logic to differentiate between Civitai and Hugging Face sources.
- Improved error handling for model downloads, providing clearer messages for authentication issues.
- Updated documentation to reflect new environment variables and usage instructions for Hugging Face integration.
- Added tests to validate the new fallback mechanism and ensure robust model resolution.
This commit is contained in:
Leonid Pershin
2026-08-21 07:34:15 +03:00
parent 7343fb0e83
commit 618e6e4806
12 changed files with 642 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/gpu_rent/capture.py
+133 -38
View File
@@ -22,6 +22,7 @@ from gpu_rent.civitai import (
)
from gpu_rent.config import Config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.huggingface import lookup_by_sha256
from gpu_rent.manifests import (
MODEL_TYPES,
extract_version_id,
@@ -49,11 +50,12 @@ def strip_git_auth(url: str) -> str:
@dataclass
class ModelCaptureItem:
kind: str
version_id: int
model_id: int
url: str
title: str = ""
rel: str = ""
version_id: int | None = None
model_id: int | None = None
source: str = "civitai" # civitai | huggingface
@dataclass
@@ -115,6 +117,7 @@ def resolve_model_item(
token: str,
api_host: str,
link_host: str,
hf_token: str | None = None,
) -> ResolveOutcome:
kind = str(raw.get("kind") or "")
if kind not in MODEL_TYPES:
@@ -133,21 +136,50 @@ def resolve_model_item(
except (TypeError, ValueError):
mid = None
def _ok(v: int, m: int, name: str = title) -> ResolveOutcome:
def _ok_civitai(v: int, m: int, name: str = title) -> ResolveOutcome:
return ResolveOutcome(
item=ModelCaptureItem(
kind=kind,
version_id=v,
model_id=m,
url=civitai_model_url(m, v, link_host),
title=name,
rel=rel,
version_id=v,
model_id=m,
source="civitai",
),
status="ok",
)
def _try_hf(sha: str, detail_prefix: str) -> ResolveOutcome | None:
"""Civitai miss → Hugging Face search by filename + LFS sha."""
try:
hit = lookup_by_sha256(
hf_token,
sha,
filename=str(raw.get("name") or rel),
)
except CloudError as exc:
return ResolveOutcome(
status="api_error",
detail=f"{detail_prefix} HF: {exc}",
)
if hit is None:
return None
return ResolveOutcome(
item=ModelCaptureItem(
kind=kind,
url=hit.url,
title=hit.title or title,
rel=rel,
version_id=None,
model_id=None,
source="huggingface",
),
status="ok",
)
if vid is not None and mid is not None:
return _ok(vid, mid)
return _ok_civitai(vid, mid)
# Partial sidecar: have version id → GET /model-versions/{id} for modelId.
if vid is not None and mid is None:
@@ -167,7 +199,7 @@ def resolve_model_item(
vid2, mid2 = version_ids_from_payload(version)
if vid2 is not None and mid2 is not None:
name = str(version.get("name") or title)
return _ok(vid2, mid2, name)
return _ok_civitai(vid2, mid2, name)
return ResolveOutcome(
status="unknown",
detail=f"{rel} version_id={vid} (нет modelId в ответе)",
@@ -183,11 +215,14 @@ def resolve_model_item(
_host, version = fetch_model_version_by_hash(token or None, api_host, str(sha))
except CloudError as exc:
msg = str(exc)
# 404 = genuinely not on Civitai; other → api_error
# 404 = genuinely not on Civitai → try HF
if "HTTP 404" in msg or msg.rstrip().endswith("404"):
hf_out = _try_hf(str(sha), f"{rel} sha={str(sha)[:12]}…")
if hf_out is not None:
return hf_out
return ResolveOutcome(
status="unknown",
detail=f"{rel} sha={str(sha)[:12]}…",
detail=f"{rel} sha={str(sha)[:12]}… (нет на Civitai/HF)",
)
return ResolveOutcome(
status="api_error",
@@ -195,12 +230,15 @@ def resolve_model_item(
)
vid2, mid2 = version_ids_from_payload(version)
if vid2 is None or mid2 is None:
hf_out = _try_hf(str(sha), f"{rel} sha={str(sha)[:12]}…")
if hf_out is not None:
return hf_out
return ResolveOutcome(
status="unknown",
detail=f"{rel} sha={str(sha)[:12]}… (пустой payload)",
)
name = str(version.get("name") or title)
return _ok(vid2, mid2, name)
return _ok_civitai(vid2, mid2, name)
def _backup(path: Path) -> None:
@@ -218,29 +256,42 @@ def _keep_model_entry(it: dict) -> bool:
return vid not in (None, "", 0, "0")
def _model_dedupe_key(kind: str, *, version_id: int | None, url: str | None) -> tuple:
u = (url or "").rstrip("/").lower()
if u and ("huggingface.co" in u or "hf.co/" in u):
return ("hf", kind, u)
if version_id is not None:
return ("civitai", kind, int(version_id))
return ("url", kind, u or "?")
def merge_models_yaml(
path: Path,
new_items: list[ModelCaptureItem],
*,
dry_run: bool,
) -> tuple[list[ModelCaptureItem], list[str]]:
"""Return (actually_new, skip_msgs). Dedupe by (kind, version_id)."""
"""Return (actually_new, skip_msgs). Dedupe by Civitai version_id or HF url."""
existing = parse_models(path) if path.is_file() else []
have: set[tuple[str, int]] = set()
have: set[tuple] = set()
for e in existing:
vid = e.version_id
if vid is None and e.url:
vid = extract_version_id(e.url)
if vid is not None:
have.add((e.kind, vid))
have.add(_model_dedupe_key(e.kind, version_id=vid, url=e.url))
added: list[ModelCaptureItem] = []
skipped: list[str] = []
seen_new: set[tuple[str, int]] = set()
seen_new: set[tuple] = set()
for item in new_items:
key = (item.kind, item.version_id)
key = _model_dedupe_key(item.kind, version_id=item.version_id, url=item.url)
if key in have or key in seen_new:
skipped.append(f"{item.kind} {item.title} modelVersionId={item.version_id}")
label = (
f"modelVersionId={item.version_id}"
if item.version_id is not None
else item.url
)
skipped.append(f"{item.kind} {item.title} {label}")
continue
seen_new.add(key)
added.append(item)
@@ -374,6 +425,7 @@ def capture_models(
link_host = cfg.civitai_api_host or "civitai.red"
token = cfg.civitai_api_token
api_host = cfg.civitai_api_host
hf_token = cfg.hf_token or None
for raw in raw_models:
if not isinstance(raw, dict):
@@ -395,18 +447,30 @@ def capture_models(
# Full sidecar / ids → no API. Partial vid → GET version. Else batch by-hash.
if vid_i is not None and mid_i is not None:
outcome = resolve_model_item(
raw, token=token, api_host=api_host, link_host=link_host
raw,
token=token,
api_host=api_host,
link_host=link_host,
hf_token=hf_token,
)
elif vid_i is not None:
outcome = resolve_model_item(
raw, token=token, api_host=api_host, link_host=link_host
raw,
token=token,
api_host=api_host,
link_host=link_host,
hf_token=hf_token,
)
elif sha:
need_hash.append(raw)
continue
else:
outcome = resolve_model_item(
raw, token=token, api_host=api_host, link_host=link_host
raw,
token=token,
api_host=api_host,
link_host=link_host,
hf_token=hf_token,
)
if outcome.status == "ok" and outcome.item is not None:
@@ -430,31 +494,62 @@ def capture_models(
report.models_api_errors.append(f"{rel} sha={sha}… (batch failed)")
need_hash = []
hf_fallback = 0
for raw in need_hash:
sha = str(raw.get("sha256") or "").strip().lower()
rel = str(raw.get("rel") or raw.get("name") or "?")
version = by_hash.get(sha)
if not version:
if version:
vid2, mid2 = version_ids_from_payload(version)
if vid2 is not None and mid2 is not None:
kind = str(raw.get("kind") or "")
if kind not in MODEL_TYPES:
continue
title = str(
version.get("name") or Path(str(raw.get("name") or rel)).stem
)
resolved.append(
ModelCaptureItem(
kind=kind,
url=civitai_model_url(mid2, vid2, link_host),
title=title,
rel=rel,
version_id=vid2,
model_id=mid2,
source="civitai",
)
)
continue
# Civitai miss → Hugging Face
try:
hit = lookup_by_sha256(
hf_token,
sha,
filename=str(raw.get("name") or rel),
)
except CloudError as exc:
report.models_api_errors.append(f"{rel} HF: {exc}")
report.models_unknown.append(f"{rel} sha={sha[:12]}…")
continue
vid2, mid2 = version_ids_from_payload(version)
if vid2 is None or mid2 is None:
report.models_unknown.append(f"{rel} sha={sha[:12]}… (пустой payload)")
continue
kind = str(raw.get("kind") or "")
if kind not in MODEL_TYPES:
continue
title = str(version.get("name") or Path(str(raw.get("name") or rel)).stem)
resolved.append(
ModelCaptureItem(
kind=kind,
version_id=vid2,
model_id=mid2,
url=civitai_model_url(mid2, vid2, link_host),
title=title,
rel=rel,
if hit:
kind = str(raw.get("kind") or "")
if kind not in MODEL_TYPES:
continue
hf_fallback += 1
resolved.append(
ModelCaptureItem(
kind=kind,
url=hit.url,
title=hit.title or Path(str(raw.get("name") or rel)).stem,
rel=rel,
source="huggingface",
)
)
)
else:
report.models_unknown.append(f"{rel} sha={sha[:12]}… (нет на Civitai/HF)")
if hf_fallback:
log(f"capture: Hugging Face fallback — {hf_fallback} файл(ов)")
added, skipped = merge_models_yaml(
cfg.models_manifest, resolved, dry_run=dry_run