mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
feat(links): support ModelScope and TensorArt as model sources
A model file could only ever be linked to huggingface.co: `set_hf_url` validated the URL with a huggingface-only regex, the agent fetched the card from a hardcoded HF URL, and the readme processor built every relative image path off `https://huggingface.co/{repo}/resolve/main`. ModelScope publishes the same model-card convention (README.md + YAML frontmatter, often carrying `base_model:` and `trigger_words:`) behind a public, key-less API, so the enrichment pipeline could already serve it - it was the plumbing that was HF-shaped, not the idea. Make the external source a first-class, provider-driven concept: - New `py/services/model_sources/` registry. A `ModelSource` owns URL recognition (lenient for stored values, strict for user input), the canonical page URL, model-card fetching, the asset base URL and the capability flags. `HuggingFaceSource` is the previous logic relocated; `ModelScopeSource` reads `/models/{o}/{n}/resolve/{master|main}/README.md` and falls back to `/api/v1/models/{o}/{n}/repo`. `TensorArtSource` is link-only on purpose: tensor.art answers plain HTTP clients with a Cloudflare challenge and its internal API (ap-east-1.tensorart.cloud / cn.tensorart.net) rejects every /v1/model/* route with "invalid authorization header", so it declares supports_enrichment=False rather than failing silently later. - Metadata gains `source_platform` + `source_url`; `hf_url` stays as a read/write alias, written only for Hugging Face, so existing sidecars, cached rows and third-party consumers keep working. Normalisation runs at the scanner, the persistent cache (both directions, plus two new columns behind an ALTER migration) and the linking handler - which is what stops a user who switches sources from leaving a stale `hf_url` on a ModelScope model. - The agent pipeline keys off the provider instead of `hf_url`: the fast-fail gate now explains *why* a model is skipped (no source / unknown source / source without a reachable card), the prompt context exposes source_url/source_id/source_label/asset_base_url while still filling the legacy hf_url/repo aliases, and the four README image extractors take a base_url (defaulting to HF) so relative paths resolve against the right site. Version grouping generalises to hf: / ms: / ta: keys. - `POST /api/lm/set-hf-url` keeps its path and its legacy payload keys but accepts `source_url`, validates against every provider and returns the platform. `GET /api/lm/model-sources` lets the UI render the supported-site list from the server. - Frontend: a `modelSourceHelpers` mirror of the registry drives the link dialog, the card/modal globe (branded "View on ModelScope/TensorArt"), the version-group key and the enrichment gate; the versions tab no longer sends ms:/ta: keys to the CivitAI API. TensorArt stays in the list because provenance is worth keeping even when the card is unreadable - the dialog says so plainly ("Sites that don't expose one (currently TensorArt) can only be linked") and the context menu disables enrichment with a matching tooltip, instead of the user getting "Unsupported URL". Verified against the real ModelScope API: jj3550945163/Krea-2-LORA returns a 1882-byte card whose frontmatter carries base_model/tags/trigger_words, and relative images resolve to .../resolve/master/.... Tests: backend 2815 passed; frontend 1130 JS + 91 Vue passed; pytest tests/i18n and a Jinja compile pass over templates/. The nine locales carry [TODO: Translate] for the new strings, completed in the next commit.
This commit is contained in:
Will Miao
@@ -0,0 +1,182 @@
|
||||
"""Tests for source-aware AI enrichment orchestration.
|
||||
|
||||
Covers the fast-fail gate (:meth:`AgentService._enrichment_skip_reason`) and
|
||||
the prompt-context builder for non-Hugging Face model sources.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.agent.agent_service import AgentService
|
||||
|
||||
|
||||
class TestEnrichmentSkipReason:
|
||||
def test_skips_when_no_source_linked(self):
|
||||
reason = AgentService._enrichment_skip_reason({})
|
||||
assert "source_url" in reason
|
||||
|
||||
def test_allows_huggingface(self):
|
||||
assert (
|
||||
AgentService._enrichment_skip_reason(
|
||||
{"hf_url": "https://huggingface.co/user/repo"}
|
||||
)
|
||||
== ""
|
||||
)
|
||||
|
||||
def test_allows_modelscope(self):
|
||||
assert (
|
||||
AgentService._enrichment_skip_reason(
|
||||
{
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/user/repo",
|
||||
}
|
||||
)
|
||||
== ""
|
||||
)
|
||||
|
||||
def test_skips_tensorart_with_reason(self):
|
||||
reason = AgentService._enrichment_skip_reason(
|
||||
{
|
||||
"source_platform": "tensorart",
|
||||
"source_url": "https://tensor.art/models/827823520299086029",
|
||||
}
|
||||
)
|
||||
assert "TensorArt" in reason
|
||||
assert "not available" in reason
|
||||
|
||||
def test_skips_unknown_platform(self):
|
||||
reason = AgentService._enrichment_skip_reason(
|
||||
{"source_platform": "somewhere", "source_url": "https://somewhere.example/m/1"}
|
||||
)
|
||||
assert "somewhere" in reason
|
||||
|
||||
|
||||
class TestBuildPromptContext:
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_card_populates_source_variables(self):
|
||||
service = AgentService()
|
||||
readme = "---\nbase_model: krea/Krea-2-Turbo\n---\n# krea\n"
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value=readme),
|
||||
) as mock_fetch,
|
||||
mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=["Krea 2 Turbo"]),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.identify_model_type",
|
||||
new=mock.AsyncMock(return_value="lora"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.settings_manager.SettingsManager.get_priority_tag_config",
|
||||
return_value={"lora": "style, subject"},
|
||||
),
|
||||
):
|
||||
context = await service._build_prompt_context(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/models/loras/krea.safetensors",
|
||||
metadata={
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
|
||||
"file_name": "krea",
|
||||
},
|
||||
registry=mock.Mock(),
|
||||
llm=mock.Mock(),
|
||||
)
|
||||
|
||||
mock_fetch.assert_awaited_once_with("jj3550945163/Krea-2-LORA")
|
||||
assert context["source_platform"] == "modelscope"
|
||||
assert context["source_id"] == "jj3550945163/Krea-2-LORA"
|
||||
assert context["source_label"] == "ModelScope"
|
||||
assert (
|
||||
context["asset_base_url"]
|
||||
== "https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master"
|
||||
)
|
||||
assert readme in context["readme_content_full"]
|
||||
# Hugging Face aliases stay empty for a non-HF source.
|
||||
assert context["hf_url"] == ""
|
||||
assert context["repo"] == "jj3550945163/Krea-2-LORA"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_huggingface_keeps_legacy_aliases(self):
|
||||
service = AgentService()
|
||||
readme = "# card\n"
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value=readme),
|
||||
) as mock_fetch,
|
||||
mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=[]),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.identify_model_type",
|
||||
new=mock.AsyncMock(return_value="lora"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.settings_manager.SettingsManager.get_priority_tag_config",
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
context = await service._build_prompt_context(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/models/loras/thing.safetensors",
|
||||
metadata={"hf_url": "https://huggingface.co/user/repo"},
|
||||
registry=mock.Mock(),
|
||||
llm=mock.Mock(),
|
||||
)
|
||||
|
||||
mock_fetch.assert_awaited_once_with("user/repo")
|
||||
assert context["source_platform"] == "huggingface"
|
||||
assert context["hf_url"] == "https://huggingface.co/user/repo"
|
||||
assert context["repo"] == "user/repo"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tensorart_never_fetches_a_card(self):
|
||||
service = AgentService()
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
|
||||
new=mock.AsyncMock(),
|
||||
) as hf_fetch,
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
|
||||
new=mock.AsyncMock(),
|
||||
) as ms_fetch,
|
||||
mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=[]),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.identify_model_type",
|
||||
new=mock.AsyncMock(return_value="lora"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.settings_manager.SettingsManager.get_priority_tag_config",
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
context = await service._build_prompt_context(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/models/loras/thing.safetensors",
|
||||
metadata={
|
||||
"source_platform": "tensorart",
|
||||
"source_url": "https://tensor.art/models/827823520299086029",
|
||||
},
|
||||
registry=mock.Mock(),
|
||||
llm=mock.Mock(),
|
||||
)
|
||||
|
||||
hf_fetch.assert_not_awaited()
|
||||
ms_fetch.assert_not_awaited()
|
||||
assert context["readme_content"] == ""
|
||||
assert context["source_platform"] == "tensorart"
|
||||
Reference in New Issue
Block a user