Add Hugging Face support and enhance model resolution logic

- 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
+62 -13
View File
@@ -79,6 +79,9 @@ def test_resolve_by_hash_404_is_unknown():
with patch(
"gpu_rent.capture.fetch_model_version_by_hash",
side_effect=CloudError("Civitai by-hash abc: HTTP 404 (хосты civitai.red)"),
), patch(
"gpu_rent.capture.lookup_by_sha256",
return_value=None,
):
out = resolve_model_item(
{
@@ -95,6 +98,40 @@ def test_resolve_by_hash_404_is_unknown():
assert out.item is None
def test_resolve_by_hash_404_hf_fallback():
from gpu_rent.huggingface import HfFileHit
with patch(
"gpu_rent.capture.fetch_model_version_by_hash",
side_effect=CloudError("Civitai by-hash abc: HTTP 404 (хосты civitai.red)"),
), patch(
"gpu_rent.capture.lookup_by_sha256",
return_value=HfFileHit(
repo_id="org/model",
filename="m.safetensors",
url="https://huggingface.co/org/model/resolve/main/m.safetensors",
sha256="a" * 64,
title="org/model",
),
):
out = resolve_model_item(
{
"kind": "lora",
"rel": "Lora/m.safetensors",
"name": "m.safetensors",
"sha256": "a" * 64,
},
token="",
api_host="civitai.red",
link_host="civitai.red",
hf_token="hf_x",
)
assert out.status == "ok"
assert out.item is not None
assert out.item.source == "huggingface"
assert "huggingface.co" in out.item.url
def test_resolve_by_hash_network_is_api_error():
with patch(
"gpu_rent.capture.fetch_model_version_by_hash",
@@ -125,17 +162,25 @@ def test_merge_models_dedupe(tmp_path: Path):
)
items = [
ModelCaptureItem(
"lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old"
kind="lora",
version_id=100,
model_id=1,
url="https://civitai.red/models/1?modelVersionId=100",
title="old",
),
ModelCaptureItem(
"lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new"
kind="lora",
version_id=200,
model_id=2,
url="https://civitai.red/models/2?modelVersionId=200",
title="new",
),
ModelCaptureItem(
"checkpoint",
300,
3,
"https://civitai.red/models/3?modelVersionId=300",
"ckpt",
kind="checkpoint",
version_id=300,
model_id=3,
url="https://civitai.red/models/3?modelVersionId=300",
title="ckpt",
),
]
added, skipped = merge_models_yaml(path, items, dry_run=False)
@@ -153,7 +198,11 @@ def test_merge_models_dry_run(tmp_path: Path):
before = path.read_text(encoding="utf-8")
items = [
ModelCaptureItem(
"lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x"
kind="lora",
version_id=1,
model_id=1,
url="https://civitai.red/models/1?modelVersionId=1",
title="x",
),
]
added, _ = merge_models_yaml(path, items, dry_run=True)
@@ -170,11 +219,11 @@ def test_merge_models_kind_scoped_dedupe(tmp_path: Path):
)
items = [
ModelCaptureItem(
"checkpoint",
100,
1,
"https://civitai.red/models/1?modelVersionId=100",
"as-ckpt",
kind="checkpoint",
version_id=100,
model_id=1,
url="https://civitai.red/models/1?modelVersionId=100",
title="as-ckpt",
),
]
added, skipped = merge_models_yaml(path, items, dry_run=False)