alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-20 18:51:26 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(download): support ModelScope repositories in the URL downloader

Browse Source 
ModelScope became a linkable source, but downloading from it was impossible:
the URL picker only recognised huggingface.co, the file listing hit a
huggingface-only endpoint, the resolve URL was hardcoded, and the default
path template always wrote into a `huggingface/` directory.

Move the download knowledge into the providers so the handlers stay generic:

- `ModelSource` gains `list_files()`, `file_download_url()`,
  `default_revision` and `default_subdir`. `HuggingFaceSource` keeps the Hub
  tree API (`/api/models/{id}/tree/{rev}`, LFS-aware sizes, `main`).
  `ModelScopeSource` uses `/api/v1/models/{id}/repo/files?Revision=master`
  — which reports real byte sizes for LFS files, so no HEAD probe is needed,
  and which only accepts `master` (an HF-imported repo still 404s on `main`)
  — and downloads through `/models/{id}/resolve/{rev}/{path}`. That URL
  redirects to a CDN target carrying a time-limited `auth_key`, so it is
  rebuilt on every request and never cached, which is also what keeps
  resumable Range requests working.
- `hf_handlers.py`/`HfHandler` become `model_source_handlers.py`/
  `ModelSourceHandler` with `list_model_source_files` and
  `download_model_source`. New routes `/api/lm/model-source-files` and
  `/api/lm/download-model-source`; the old `/api/lm/hf-repo-files` and
  `/api/lm/download-hf-model` paths stay as aliases, and a payload without
  `platform` still means Hugging Face, so existing callers are unaffected.
- A downloaded sidecar now records `source_platform` + `source_url` (with the
  `hf_url` alias only for Hugging Face) instead of always writing `hf_url`,
  and `use_default_paths` files ModelScope downloads under
  `modelscope/<owner>/<repo>`. The now-unused shared HF aiohttp session and
  its shutdown hook are gone; providers open short-lived sessions.
- Frontend: `detectUrlType` returns the platform-neutral
  `model-source-repo` / `model-source-file` plus an explicit `platform`, the
  DownloadManager's `hf*` state and methods are renamed to `source*`, every
  `source === 'huggingface'` check becomes `isExternalModelSource()`, and
  batch groups are keyed by `platform:repo` so the same `owner/name` on two
  sites renders as two groups. A bare `owner/name` still means Hugging Face.
- `is_valid_source_id()` centralises repo-id validation (exactly
  `owner/name`, no traversal, no leading dot). This also fixes the old HF
  download check that rejected any dot in the name, i.e. legitimate repos
  such as `black-forest-labs/FLUX.1-dev`.

Verified against the live APIs: the example repo lists 8 weight files with
correct sizes, and a ranged GET of the built resolve URL returns 206 after
following the redirect to the CDN. Backend 2853 passed; frontend 1143 JS +
91 Vue passed. The nine locales carry the refreshed download copy in the
next commit.
This commit is contained in:
Will Miao
2026-09-14 07:42:51 +08:00
parent b9bf006998
commit 38d4c59b4c
22 changed files with 1953 additions and 667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/frontend/managers/downloadManager.batchSummary.test.js
+12 -8
View File
@@ -26,6 +26,7 @@ const {
},
},
downloadModel: vi.fn(),
downloadModelSource: vi.fn(),
downloadHfModel: vi.fn(),
cancelDownload: vi.fn(),
getPageState: vi.fn(() => ({})),
@@ -158,7 +159,7 @@ describe('DownloadManager batch download summary flow', () => {
// Reset the shared mocks so mockResolvedValueOnce queues and call
// history never leak between tests.
mockApiClient.downloadModel.mockReset();
mockApiClient.downloadHfModel.mockReset();
mockApiClient.downloadModelSource.mockReset();
mockApiClient.cancelDownload.mockReset();
showToastMock.mockClear();
showDownloadBatchSummaryMock.mockClear();
@@ -406,14 +407,15 @@ describe('DownloadManager batch download summary flow', () => {
expect(showToastMock).toHaveBeenCalledWith('toast.loras.downloadCompleted', {}, 'success');
});
it('shows a summary for HF partial failure and retries only the failed files', async () => {
manager.hfRepoId = 'user/repo';
manager.hfSelectedFiles = ['a.safetensors', 'b.safetensors'];
mockApiClient.downloadHfModel
it('shows a summary for external repo partial failure and retries only the failed files', async () => {
manager.sourcePlatform = 'huggingface';
manager.sourceRepoId = 'user/repo';
manager.sourceSelectedFiles = ['a.safetensors', 'b.safetensors'];
mockApiClient.downloadModelSource
.mockResolvedValueOnce({ success: true })
.mockResolvedValueOnce({ success: false, error: 'denied' });
const result = await manager._downloadHfSingle({ modelRoot: '/m', useDefaultPaths: true });
const result = await manager._downloadExternalRepoFiles({ modelRoot: '/m', useDefaultPaths: true });
expect(result).toBe(false);
expect(showDownloadBatchSummaryMock).toHaveBeenCalledTimes(1);
@@ -428,7 +430,9 @@ describe('DownloadManager batch download summary flow', () => {
await summary.onRetry();
expect(mockApiClient.downloadHfModel).toHaveBeenCalledTimes(3);
expect(mockApiClient.downloadHfModel.mock.calls[2][0].filename).toBe('b.safetensors');
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(3);
const retryArgs = mockApiClient.downloadModelSource.mock.calls[2][0];
expect(retryArgs.filename).toBe('b.safetensors');
expect(retryArgs.platform).toBe('huggingface');
});
});
tests/frontend/managers/downloadManager.modelSource.test.js
+269
View File
@@ -0,0 +1,269 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
const {
DOWNLOAD_MANAGER_MODULE,
MODAL_MANAGER_MODULE,
UI_HELPERS_MODULE,
STATE_MODULE,
LOADING_MANAGER_MODULE,
API_FACTORY_MODULE,
STORAGE_HELPERS_MODULE,
FOLDER_TREE_MANAGER_MODULE,
I18N_HELPERS_MODULE,
SUMMARY_MODULE,
mockApiClient,
mockLoadingManager,
showDownloadBatchSummaryMock,
} = vi.hoisted(() => {
const mockApiClient = {
apiConfig: { config: { displayName: 'LoRA', singularName: 'lora' } },
downloadModel: vi.fn(),
downloadModelSource: vi.fn(),
fetchModelSourceFiles: vi.fn(),
cancelDownload: vi.fn(),
getPageState: vi.fn(() => ({})),
};
const mockLoadingManager = {
showSimpleLoading: vi.fn(),
hide: vi.fn(),
restoreProgressBar: vi.fn(),
showDownloadProgress: vi.fn(() => vi.fn()),
setStatus: vi.fn(),
showCancelButton: vi.fn(),
};
return {
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname,
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname,
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
SUMMARY_MODULE: new URL('../../../static/js/components/DownloadBatchSummaryModal.js', import.meta.url).pathname,
mockApiClient,
mockLoadingManager,
showDownloadBatchSummaryMock: vi.fn(),
};
});
vi.mock(MODAL_MANAGER_MODULE, () => ({
modalManager: { showModal: vi.fn(), closeModal: vi.fn() },
}));
vi.mock(UI_HELPERS_MODULE, () => ({ showToast: vi.fn() }));
vi.mock(STATE_MODULE, () => ({
state: { global: { settings: {} }, loadingManager: mockLoadingManager },
}));
vi.mock(LOADING_MANAGER_MODULE, () => ({
LoadingManager: vi.fn(() => mockLoadingManager),
}));
vi.mock(API_FACTORY_MODULE, () => ({
getModelApiClient: vi.fn(() => mockApiClient),
resetAndReload: vi.fn(),
}));
vi.mock(STORAGE_HELPERS_MODULE, () => ({
getStorageItem: vi.fn((_key, defaultValue) => defaultValue),
setStorageItem: vi.fn(),
}));
vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({
FolderTreeManager: vi.fn(() => ({ clearSelection: vi.fn(), init: vi.fn() })),
}));
vi.mock(I18N_HELPERS_MODULE, () => ({
translate: vi.fn((_, __, fallback) => fallback ?? ''),
}));
vi.mock(SUMMARY_MODULE, () => ({
showDownloadBatchSummary: showDownloadBatchSummaryMock,
}));
class FakeWebSocket {
constructor(url) {
this.url = url;
this.onopen = null;
this.onmessage = null;
this.onerror = null;
this.close = vi.fn();
queueMicrotask(() => {
if (this.onopen) this.onopen();
});
}
}
const MS_REPO_URL = 'https://modelscope.cn/models/jj3550945163/Krea-2-LORA';
describe('DownloadManager external model source downloads', () => {
let DownloadManager;
let manager;
beforeEach(async () => {
document.body.innerHTML = '';
vi.stubGlobal('WebSocket', FakeWebSocket);
mockApiClient.downloadModelSource.mockReset();
mockApiClient.fetchModelSourceFiles.mockReset();
mockLoadingManager.showSimpleLoading.mockReset();
({ DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE));
manager = new DownloadManager();
manager.apiClient = mockApiClient;
manager.showBatchPreviewStep = vi.fn();
manager.proceedToLocation = vi.fn();
});
it('loads a ModelScope repo as batch items on the master revision', async () => {
mockApiClient.fetchModelSourceFiles.mockResolvedValue([
{ filename: 'a.safetensors', size: 10 },
{ filename: 'sub/b.safetensors', size: 20 },
]);
const errorElement = { textContent: '' };
await manager._validateAndFetchExternalRepo([MS_REPO_URL], errorElement);
expect(mockApiClient.fetchModelSourceFiles).toHaveBeenCalledWith(
'jj3550945163/Krea-2-LORA',
'modelscope',
'master'
);
expect(errorElement.textContent).toBe('');
expect(manager.source).toBe('modelscope');
expect(manager.isBatchMode).toBe(true);
expect(manager.batchModels).toHaveLength(2);
expect(manager.batchModels[0]).toMatchObject({
source: 'modelscope',
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
revision: 'master',
filename: 'a.safetensors',
fileSizeBytes: 10,
displayName: 'a.safetensors',
});
expect(manager.showBatchPreviewStep).toHaveBeenCalled();
});
it('keeps Hugging Face on its own revision', async () => {
mockApiClient.fetchModelSourceFiles.mockResolvedValue([
{ filename: 'a.safetensors', size: 10 },
]);
await manager._validateAndFetchExternalRepo(
['https://huggingface.co/user/repo'],
{ textContent: '' }
);
expect(mockApiClient.fetchModelSourceFiles).toHaveBeenCalledWith(
'user/repo',
'huggingface',
'main'
);
expect(manager.batchModels[0].revision).toBe('main');
});
it('surfaces a listing failure on the URL field', async () => {
mockApiClient.fetchModelSourceFiles.mockRejectedValue(new Error('Repository not found'));
const errorElement = { textContent: '' };
await manager._validateAndFetchExternalRepo([MS_REPO_URL], errorElement);
expect(errorElement.textContent).toBe('Repository not found');
expect(manager.showBatchPreviewStep).not.toHaveBeenCalled();
});
it('skips file selection for a direct ModelScope file URL', async () => {
await manager._validateAndFetchExternalRepo(
[`${MS_REPO_URL}/resolve/master/Krea-2-LORA_c1-st1000.safetensors`],
{ textContent: '' }
);
expect(manager.isBatchMode).toBe(false);
expect(manager.sourcePlatform).toBe('modelscope');
expect(manager.sourceRepoId).toBe('jj3550945163/Krea-2-LORA');
expect(manager.sourceSelectedFiles).toEqual(['Krea-2-LORA_c1-st1000.safetensors']);
expect(manager.proceedToLocation).toHaveBeenCalled();
});
it('downloads a single ModelScope file through the generic endpoint', async () => {
mockApiClient.downloadModelSource.mockResolvedValue({ success: true });
manager.sourcePlatform = 'modelscope';
manager.sourceRepoId = 'jj3550945163/Krea-2-LORA';
manager.sourceSelectedFiles = ['Krea-2-LORA_c1-st1000.safetensors'];
await manager._downloadExternalRepoFiles({
modelRoot: '/models/loras',
targetFolder: '',
useDefaultPaths: true,
});
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(1);
expect(mockApiClient.downloadModelSource.mock.calls[0][0]).toMatchObject({
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
filename: 'Krea-2-LORA_c1-st1000.safetensors',
revision: 'master',
});
});
it('carries the platform through a batch download', async () => {
mockApiClient.downloadModelSource.mockResolvedValue({ success: true });
manager.showBatchPreviewStep = vi.fn();
await manager.executeBatchDownload(
[
{
source: 'modelscope',
platform: 'modelscope',
repo: 'u/r',
filename: 'f.safetensors',
revision: 'master',
displayName: 'f.safetensors',
checked: true,
},
],
{ modelRoot: '/models/loras', targetFolder: '', useDefaultPaths: true }
);
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(1);
expect(mockApiClient.downloadModelSource.mock.calls[0][0]).toMatchObject({
platform: 'modelscope',
repo: 'u/r',
filename: 'f.safetensors',
revision: 'master',
});
expect(mockApiClient.downloadModel).not.toHaveBeenCalled();
});
it('links failures to the ModelScope file page', async () => {
expect(
manager._buildSingleItemUrl({
source: 'modelscope',
repo: 'u/r',
filename: 'sub/f.safetensors',
})
).toBe('https://modelscope.cn/models/u/r/file/view/master/sub/f.safetensors');
expect(
manager._buildSingleItemUrl({
source: 'huggingface',
repo: 'u/r',
filename: 'f.safetensors',
})
).toBe('https://huggingface.co/u/r/blob/main/f.safetensors');
});
it('groups the same repo name on two platforms separately', () => {
const hf = { source: 'huggingface', repo: 'u/r' };
const ms = { source: 'modelscope', repo: 'u/r' };
expect(manager._externalGroupKey(hf)).toBe('huggingface:u/r');
expect(manager._externalGroupKey(ms)).toBe('modelscope:u/r');
expect(manager._externalGroupKey(hf)).not.toBe(manager._externalGroupKey(ms));
});
});
tests/frontend/utils/hfUrlDetection.test.js → tests/frontend/utils/modelSourceUrlDetection.test.js
+61 -8
View File
@@ -1,14 +1,15 @@
import { describe, it, expect } from 'vitest';
import { DownloadManager } from '../../../static/js/managers/DownloadManager.js';
describe('DownloadManager.detectUrlType — HF URL detection', () => {
describe('DownloadManager.detectUrlType — external model source URLs', () => {
it('detects HF resolve URL with file', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency/resolve/main/Flux2-Klein-9B-consistency-V2.safetensors'
);
expect(result).toEqual({
type: 'hf-resolve',
type: 'model-source-file',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
revision: 'main',
filename: 'Flux2-Klein-9B-consistency-V2.safetensors',
@@ -20,7 +21,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/user/repo/resolve/main/subdir/model.safetensors'
);
expect(result).toEqual({
type: 'hf-resolve',
type: 'model-source-file',
platform: 'huggingface',
repo: 'user/repo',
revision: 'main',
filename: 'subdir/model.safetensors',
@@ -32,7 +34,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency'
);
expect(result).toEqual({
type: 'hf-repo',
type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
});
});
@@ -40,7 +43,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
it('detects HF repo URL (bare user/repo)', () => {
const result = DownloadManager.detectUrlType('dx8152/Flux2-Klein-9B-Consistency');
expect(result).toEqual({
type: 'hf-repo',
type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
});
});
@@ -50,7 +54,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/user/repo/'
);
expect(result).toEqual({
type: 'hf-repo',
type: 'model-source-repo',
platform: 'huggingface',
repo: 'user/repo',
});
});
@@ -60,7 +65,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors'
);
expect(result).toEqual({
type: 'hf-resolve',
type: 'model-source-file',
platform: 'huggingface',
repo: 'Comfy-Org/z_image_turbo',
revision: 'main',
filename: 'split_files/diffusion_models/z_image_turbo_bf16.safetensors',
@@ -115,7 +121,7 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/user/repo/resolve/main/file.safetensors'
);
expect(result?.type).toBe('hf-resolve');
expect(result?.type).toBe('model-source-file');
});
it('prefers CivitAI over HF when both match', () => {
@@ -126,4 +132,51 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
);
expect(result?.type).toBe('civitai');
});
it('detects a ModelScope repo URL', () => {
const result = DownloadManager.detectUrlType(
'https://modelscope.cn/models/jj3550945163/Krea-2-LORA'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
});
});
it('detects a ModelScope file URL with revision and subdirectory', () => {
const result = DownloadManager.detectUrlType(
'https://modelscope.cn/models/AI-ModelScope/stable-diffusion-v1-5/resolve/master/vae/diffusion_pytorch_model.bin'
);
expect(result).toEqual({
type: 'model-source-file',
platform: 'modelscope',
repo: 'AI-ModelScope/stable-diffusion-v1-5',
revision: 'master',
filename: 'vae/diffusion_pytorch_model.bin',
});
});
it('detects a ModelScope view sub-page as a repo URL', () => {
const result = DownloadManager.detectUrlType(
'https://www.modelscope.cn/models/user/repo/summary'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'modelscope',
repo: 'user/repo',
});
});
it('does not treat a bare owner/name as ModelScope', () => {
// The shorthand has always meant Hugging Face; ModelScope needs its host.
const result = DownloadManager.detectUrlType('user/repo');
expect(result.platform).toBe('huggingface');
});
it('rejects path traversal in either platform', () => {
expect(
DownloadManager.detectUrlType('https://modelscope.cn/models/../etc/passwd')
).toBeNull();
});
});
tests/routes/test_hf_handlers.py
-308
View File
@@ -1,308 +0,0 @@
"""Tests for the HuggingFace link handler (``set_hf_url``).
Regression coverage for issue #1094: linking a model to HuggingFace must not
clear its CivitAI provenance or metadata, so both "View on CivitAI" and
"View on Hugging Face" can coexist.
"""
from __future__ import annotations
import json
import os
from typing import Any
from unittest.mock import AsyncMock
import pytest
from py.routes.handlers import hf_handlers
from py.routes.handlers.hf_handlers import HfHandler
from py.utils.metadata_manager import MetadataManager
def _json_payload(response) -> dict[str, Any]:
assert response.text is not None
return json.loads(response.text)
class FakeRequest:
def __init__(self, *, json_data=None):
self._json_data = json_data or {}
async def json(self):
return self._json_data
def _sidecar_path(model_path) -> str:
return f"{os.path.splitext(str(model_path))[0]}.metadata.json"
@pytest.fixture
def hf_env(tmp_path, monkeypatch):
"""Point HF linking at *tmp_path* and stub the scanner cache write."""
monkeypatch.setattr(hf_handlers, "_find_matching_root", lambda _dir: str(tmp_path))
cache_write = AsyncMock()
monkeypatch.setattr(hf_handlers, "_add_to_scanner_cache", cache_write)
return {"root": tmp_path, "cache_write": cache_write}
async def _write_model(model_path, payload: dict[str, Any]) -> None:
model_path.write_bytes(b"x" * 32)
await MetadataManager.save_metadata(str(model_path), payload)
@pytest.mark.asyncio
async def test_set_hf_url_keeps_civitai_metadata_and_provenance(tmp_path, hf_env):
model_path = tmp_path / "civitai_model.safetensors"
await _write_model(
model_path,
{
"file_name": "civitai_model",
"model_name": "CivitAI Model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "a" * 64,
"base_model": "SDXL 1.0",
"preview_url": "",
"from_civitai": True,
"civitai": {"id": 111, "modelId": 222, "name": "v1", "trainedWords": []},
},
)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
assert _json_payload(response)["success"] is True
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
# Linking HF must not erase the model's CivitAI provenance or data.
assert saved["from_civitai"] is True
assert saved["civitai"]["modelId"] == 222
assert saved["civitai"]["id"] == 111
hf_env["cache_write"].assert_awaited_once()
cached_metadata = hf_env["cache_write"].await_args.args[1]
assert cached_metadata["hf_url"] == "https://huggingface.co/user/repo"
assert cached_metadata["from_civitai"] is True
assert cached_metadata["civitai"]["modelId"] == 222
@pytest.mark.asyncio
async def test_set_hf_url_does_not_force_from_civitai_false(tmp_path, hf_env):
"""A model without CivitAI data keeps its existing provenance flag."""
model_path = tmp_path / "hf_only.safetensors"
await _write_model(
model_path,
{
"file_name": "hf_only",
"model_name": "HF Only",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "b" * 64,
"base_model": "Unknown",
"preview_url": "",
"from_civitai": True,
},
)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
assert saved["from_civitai"] is True
@pytest.mark.asyncio
async def test_set_hf_url_rejects_non_repo_url(tmp_path, hf_env):
model_path = tmp_path / "model.safetensors"
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "c" * 64,
"base_model": "Unknown",
"preview_url": "",
},
)
response = await HfHandler().set_hf_url(
FakeRequest(json_data={"file_path": str(model_path), "hf_url": "https://example.com/x"})
)
assert response.status == 400
payload = _json_payload(response)
assert payload["success"] is False
hf_env["cache_write"].assert_not_awaited()
# ---------------------------------------------------------------------------
# Multi-source linking (ModelScope / TensorArt)
# ---------------------------------------------------------------------------
async def _write_plain_model(model_path, sha: str = "d" * 64) -> None:
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": sha,
"base_model": "Unknown",
"preview_url": "",
},
)
@pytest.mark.asyncio
async def test_set_hf_url_accepts_modelscope_and_stores_source_fields(tmp_path, hf_env):
model_path = tmp_path / "ms_model.safetensors"
await _write_plain_model(model_path)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "modelscope"
assert payload["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
# No stale Hugging Face alias for a ModelScope model.
assert saved.get("hf_url", "") == ""
cached_metadata = hf_env["cache_write"].await_args.args[1]
assert cached_metadata["source_platform"] == "modelscope"
@pytest.mark.asyncio
async def test_set_hf_url_accepts_tensorart_url(tmp_path, hf_env):
model_path = tmp_path / "ta_model.safetensors"
await _write_plain_model(model_path, sha="e" * 64)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": (
"https://tensor.art/models/827823520299086029/"
"Vivid-Impressions-Storybook-Sstyle-V1.0"
),
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "tensorart"
# The canonical page URL is stored, without the slug.
assert payload["source_url"] == "https://tensor.art/models/827823520299086029"
@pytest.mark.asyncio
async def test_set_hf_url_canonicalises_modelscope_subpage(tmp_path, hf_env):
model_path = tmp_path / "ms_sub.safetensors"
await _write_plain_model(model_path, sha="f" * 64)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo/summary",
}
)
)
assert response.status == 200
assert _json_payload(response)["source_url"] == "https://modelscope.cn/models/user/repo"
@pytest.mark.asyncio
async def test_set_hf_url_is_idempotent_for_modelscope(tmp_path, hf_env):
model_path = tmp_path / "ms_twice.safetensors"
await _write_plain_model(model_path, sha="1" * 64)
request = FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
await HfHandler().set_hf_url(request)
await HfHandler().set_hf_url(request)
# The second call short-circuits without rewriting the cache entry.
assert hf_env["cache_write"].await_count == 1
@pytest.mark.asyncio
async def test_set_hf_url_switching_source_clears_hf_alias(tmp_path, hf_env):
model_path = tmp_path / "switch.safetensors"
await _write_plain_model(model_path, sha="2" * 64)
await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://huggingface.co/user/repo",
}
)
)
await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
)
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved.get("hf_url", "") == ""
@pytest.mark.asyncio
async def test_get_model_sources_lists_capabilities():
response = await HfHandler().get_model_sources(FakeRequest())
sources = _json_payload(response)
by_platform = {s["platform"]: s for s in sources}
assert set(by_platform) == {"huggingface", "modelscope", "tensorart"}
assert by_platform["huggingface"]["supports_enrichment"] is True
assert by_platform["modelscope"]["supports_enrichment"] is True
# TensorArt is link-only: no accessible model card for the backend.
assert by_platform["tensorart"]["supports_enrichment"] is False
assert by_platform["modelscope"]["supports_download"] is False
assert all(s["example_url"] for s in sources)
tests/routes/test_model_source_handlers.py
+635
View File
@@ -0,0 +1,635 @@
"""Tests for the external model-source handlers.
Covers linking (``set_hf_url``), file listing and downloads across the
registered platforms (Hugging Face / ModelScope).
Regression coverage for issue #1094: linking a model to HuggingFace must not
clear its CivitAI provenance or metadata, so both "View on CivitAI" and
"View on Hugging Face" can coexist.
"""
from __future__ import annotations
import json
import os
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock
import pytest
from py.routes.handlers import model_source_handlers
from py.routes.handlers.model_source_handlers import ModelSourceHandler
from py.services.model_sources import ModelSourceError, SourceRef
from py.services.service_registry import ServiceRegistry
from py.utils.models import LoraMetadata
from py.utils.metadata_manager import MetadataManager
def _json_payload(response) -> dict[str, Any]:
assert response.text is not None
return json.loads(response.text)
class FakeRequest:
def __init__(self, *, json_data=None, query=None):
self._json_data = json_data or {}
self.query = query or {}
async def json(self):
return self._json_data
def _sidecar_path(model_path) -> str:
return f"{os.path.splitext(str(model_path))[0]}.metadata.json"
@pytest.fixture
def source_env(tmp_path, monkeypatch):
"""Point HF linking at *tmp_path* and stub the scanner cache write."""
monkeypatch.setattr(model_source_handlers, "_find_matching_root", lambda _dir: str(tmp_path))
cache_write = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_add_to_scanner_cache", cache_write)
return {"root": tmp_path, "cache_write": cache_write}
async def _write_model(model_path, payload: dict[str, Any]) -> None:
model_path.write_bytes(b"x" * 32)
await MetadataManager.save_metadata(str(model_path), payload)
@pytest.mark.asyncio
async def test_set_hf_url_keeps_civitai_metadata_and_provenance(tmp_path, source_env):
model_path = tmp_path / "civitai_model.safetensors"
await _write_model(
model_path,
{
"file_name": "civitai_model",
"model_name": "CivitAI Model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "a" * 64,
"base_model": "SDXL 1.0",
"preview_url": "",
"from_civitai": True,
"civitai": {"id": 111, "modelId": 222, "name": "v1", "trainedWords": []},
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
assert _json_payload(response)["success"] is True
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
# Linking HF must not erase the model's CivitAI provenance or data.
assert saved["from_civitai"] is True
assert saved["civitai"]["modelId"] == 222
assert saved["civitai"]["id"] == 111
source_env["cache_write"].assert_awaited_once()
cached_metadata = source_env["cache_write"].await_args.args[1]
assert cached_metadata["hf_url"] == "https://huggingface.co/user/repo"
assert cached_metadata["from_civitai"] is True
assert cached_metadata["civitai"]["modelId"] == 222
@pytest.mark.asyncio
async def test_set_hf_url_does_not_force_from_civitai_false(tmp_path, source_env):
"""A model without CivitAI data keeps its existing provenance flag."""
model_path = tmp_path / "hf_only.safetensors"
await _write_model(
model_path,
{
"file_name": "hf_only",
"model_name": "HF Only",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "b" * 64,
"base_model": "Unknown",
"preview_url": "",
"from_civitai": True,
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
assert saved["from_civitai"] is True
@pytest.mark.asyncio
async def test_set_hf_url_rejects_non_repo_url(tmp_path, source_env):
model_path = tmp_path / "model.safetensors"
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "c" * 64,
"base_model": "Unknown",
"preview_url": "",
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(json_data={"file_path": str(model_path), "hf_url": "https://example.com/x"})
)
assert response.status == 400
payload = _json_payload(response)
assert payload["success"] is False
source_env["cache_write"].assert_not_awaited()
# ---------------------------------------------------------------------------
# Multi-source linking (ModelScope / TensorArt)
# ---------------------------------------------------------------------------
async def _write_plain_model(model_path, sha: str = "d" * 64) -> None:
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": sha,
"base_model": "Unknown",
"preview_url": "",
},
)
@pytest.mark.asyncio
async def test_set_hf_url_accepts_modelscope_and_stores_source_fields(tmp_path, source_env):
model_path = tmp_path / "ms_model.safetensors"
await _write_plain_model(model_path)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "modelscope"
assert payload["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
# No stale Hugging Face alias for a ModelScope model.
assert saved.get("hf_url", "") == ""
cached_metadata = source_env["cache_write"].await_args.args[1]
assert cached_metadata["source_platform"] == "modelscope"
@pytest.mark.asyncio
async def test_set_hf_url_accepts_tensorart_url(tmp_path, source_env):
model_path = tmp_path / "ta_model.safetensors"
await _write_plain_model(model_path, sha="e" * 64)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": (
"https://tensor.art/models/827823520299086029/"
"Vivid-Impressions-Storybook-Sstyle-V1.0"
),
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "tensorart"
# The canonical page URL is stored, without the slug.
assert payload["source_url"] == "https://tensor.art/models/827823520299086029"
@pytest.mark.asyncio
async def test_set_hf_url_canonicalises_modelscope_subpage(tmp_path, source_env):
model_path = tmp_path / "ms_sub.safetensors"
await _write_plain_model(model_path, sha="f" * 64)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo/summary",
}
)
)
assert response.status == 200
assert _json_payload(response)["source_url"] == "https://modelscope.cn/models/user/repo"
@pytest.mark.asyncio
async def test_set_hf_url_is_idempotent_for_modelscope(tmp_path, source_env):
model_path = tmp_path / "ms_twice.safetensors"
await _write_plain_model(model_path, sha="1" * 64)
request = FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
await ModelSourceHandler().set_hf_url(request)
await ModelSourceHandler().set_hf_url(request)
# The second call short-circuits without rewriting the cache entry.
assert source_env["cache_write"].await_count == 1
@pytest.mark.asyncio
async def test_set_hf_url_switching_source_clears_hf_alias(tmp_path, source_env):
model_path = tmp_path / "switch.safetensors"
await _write_plain_model(model_path, sha="2" * 64)
await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://huggingface.co/user/repo",
}
)
)
await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
)
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved.get("hf_url", "") == ""
@pytest.mark.asyncio
async def test_get_model_sources_lists_capabilities():
response = await ModelSourceHandler().get_model_sources(FakeRequest())
sources = _json_payload(response)
by_platform = {s["platform"]: s for s in sources}
assert set(by_platform) == {"huggingface", "modelscope", "tensorart"}
assert by_platform["huggingface"]["supports_enrichment"] is True
assert by_platform["modelscope"]["supports_enrichment"] is True
# TensorArt is link-only: no accessible model card for the backend.
assert by_platform["tensorart"]["supports_enrichment"] is False
assert by_platform["modelscope"]["supports_download"] is True
assert by_platform["modelscope"]["default_revision"] == "master"
assert by_platform["tensorart"]["supports_download"] is False
assert all(s["example_url"] for s in sources)
# ---------------------------------------------------------------------------
# File listing
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_list_model_source_files_returns_provider_result(monkeypatch):
captured: dict = {}
async def fake_list_files(self, source_id, revision=""):
captured["source_id"] = source_id
captured["revision"] = revision
return [{"filename": "a.safetensors", "size": 10}]
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(
query={
"platform": "modelscope",
"repo": "jj3550945163/Krea-2-LORA",
"revision": "v1",
}
)
)
assert response.status == 200
assert _json_payload(response) == [{"filename": "a.safetensors", "size": 10}]
assert captured == {"source_id": "jj3550945163/Krea-2-LORA", "revision": "v1"}
@pytest.mark.asyncio
async def test_list_model_source_files_rejects_link_only_platform():
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "tensorart", "repo": "u/r"})
)
assert response.status == 400
assert "does not support downloads" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_rejects_unsafe_repo():
for repo in ("noslash", "../etc/passwd", "u/.."):
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": repo})
)
assert response.status == 400, repo
assert "repo" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_maps_missing_repo_to_404(monkeypatch):
async def fake_list_files(self, source_id, revision=""):
raise ModelSourceError(f"Repository '{source_id}' not found", status=404)
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": "u/r"})
)
assert response.status == 404
assert "not found" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_maps_transport_failure_to_502(monkeypatch):
async def fake_list_files(self, source_id, revision=""):
raise ModelSourceError("upstream exploded", status=502)
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": "u/r"})
)
assert response.status == 502
# ---------------------------------------------------------------------------
# Downloads
# ---------------------------------------------------------------------------
def _stub_download_backend(monkeypatch) -> dict:
"""Replace the downloader/settings plumbing with a recording stub."""
captured: dict = {}
async def fake_download_file(**kwargs):
captured.update(kwargs)
return True, kwargs["save_path"]
class _Downloader:
download_file = staticmethod(fake_download_file)
async def fake_get_downloader():
return _Downloader()
class _Settings:
def get(self, key, default=None):
return default
monkeypatch.setattr(model_source_handlers, "get_downloader", fake_get_downloader)
monkeypatch.setattr(
model_source_handlers, "get_settings_manager", lambda: _Settings()
)
return captured
@pytest.mark.asyncio
async def test_download_model_source_modelscope_uses_resolve_url(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "jj3550945163/Krea-2-LORA",
"filename": "Krea-2-LORA_c1-st1000.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["url"] == (
"https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master/"
"Krea-2-LORA_c1-st1000.safetensors"
)
assert captured["save_path"] == str(tmp_path / "Krea-2-LORA_c1-st1000.safetensors")
ref = saved.await_args.args[1]
assert ref.platform == "modelscope"
assert ref.source_id == "jj3550945163/Krea-2-LORA"
assert ref.url == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
@pytest.mark.asyncio
async def test_download_model_source_modelscope_default_paths(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "owner/name",
"filename": "nested/model.safetensors",
"model_root": str(tmp_path),
"use_default_paths": True,
}
)
)
assert response.status == 200
# The site gets its own sub-directory, mirroring `huggingface/<owner>/<repo>`.
assert captured["save_path"] == str(
tmp_path / "modelscope" / "owner" / "name" / "model.safetensors"
)
@pytest.mark.asyncio
async def test_download_model_source_defaults_to_huggingface(tmp_path, monkeypatch):
"""The legacy /api/lm/download-hf-model payload has no `platform` key."""
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"repo": "user/repo",
"filename": "f.safetensors",
"revision": "main",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["url"] == (
"https://huggingface.co/user/repo/resolve/main/f.safetensors"
)
@pytest.mark.asyncio
async def test_download_model_source_rejects_link_only_platform(tmp_path):
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "tensorart",
"repo": "u/r",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 400
assert "does not support downloads" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_download_model_source_rejects_unsafe_input(tmp_path, monkeypatch):
_stub_download_backend(monkeypatch)
cases = [
({"repo": "noslash", "filename": "f.safetensors"}, "repo format"),
({"repo": "u/r", "filename": "../../etc/passwd"}, "Invalid filename"),
(
{"repo": "u/r", "filename": "f.safetensors", "relative_path": "/abs"},
"relative_path must not be absolute",
),
(
{"repo": "u/r", "filename": "f.safetensors", "relative_path": "../up"},
"Invalid relative_path",
),
]
for extra, expected in cases:
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"model_root": str(tmp_path),
**extra,
}
)
)
assert response.status == 400, extra
assert expected in _json_payload(response)["error"], extra
@pytest.mark.asyncio
async def test_download_model_source_skips_existing_file(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
(tmp_path / "f.safetensors").write_bytes(b"already here")
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "u/r",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert "already exists" in _json_payload(response)["message"]
assert captured == {}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("platform", "url", "expect_hf_alias"),
[
("modelscope", "https://modelscope.cn/models/u/r", False),
("huggingface", "https://huggingface.co/u/r", True),
],
)
async def test_save_source_metadata_writes_platform_fields(
tmp_path, monkeypatch, platform, url, expect_hf_alias
):
"""A download's sidecar must record its own platform (and no stale HF alias)."""
model_path = tmp_path / "downloaded.safetensors"
model_path.write_bytes(b"x" * 32)
metadata = LoraMetadata(
file_name="downloaded",
model_name="Downloaded",
file_path=str(model_path),
size=32,
modified=1.0,
sha256="a" * 64,
base_model="SDXL 1.0",
preview_url="",
)
monkeypatch.setattr(
model_source_handlers.MetadataManager,
"create_default_metadata",
AsyncMock(return_value=metadata),
)
scanner = SimpleNamespace(add_model_to_cache=AsyncMock())
monkeypatch.setattr(
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=scanner)
)
monkeypatch.setattr(
model_source_handlers, "_infer_model_type", lambda _root: (LoraMetadata, "get_lora_scanner")
)
ref = SourceRef(platform=platform, source_id="u/r", url=url)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == platform
assert saved["source_url"] == url
assert bool(saved.get("hf_url", "")) is expect_hf_alias
cached = scanner.add_model_to_cache.await_args.args[0]
assert cached["source_platform"] == platform
assert cached["source_url"] == url
tests/services/test_model_sources.py
+184 -2
View File
@@ -14,10 +14,14 @@ from py.services.model_sources import (
ModelScopeSource,
TensorArtSource,
detect_source,
downloadable_sources,
get_download_source,
get_source,
get_source_platform,
has_external_source,
is_valid_source_id,
list_sources,
ModelSourceError,
normalize_metadata_source,
resolve_source_ref,
source_group_key,
@@ -127,11 +131,15 @@ class TestCapabilities:
source = get_source("huggingface")
assert source.supports_enrichment is True
assert source.supports_download is True
assert source.default_revision == "main"
assert source.default_subdir == "huggingface"
def test_modelscope_supports_enrichment_but_not_download(self):
def test_modelscope_supports_enrichment_and_download(self):
source = get_source("modelscope")
assert source.supports_enrichment is True
assert source.supports_download is False
assert source.supports_download is True
assert source.default_revision == "master"
assert source.default_subdir == "modelscope"
def test_tensorart_is_link_only(self):
source = get_source("tensorart")
@@ -322,3 +330,177 @@ class TestAssetBaseUrl:
ModelScopeSource().asset_base_url("u/r")
== "https://modelscope.cn/models/u/r/resolve/master"
)
# ---------------------------------------------------------------------------
# Download support
# ---------------------------------------------------------------------------
class TestListFiles:
@pytest.mark.asyncio
async def test_huggingface_reads_tree_api_with_lfs_sizes(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, [
{"path": "README.md", "size": 120},
{"path": "a/model.safetensors", "size": 300},
{"path": "b.safetensors", "size": 0, "lfs": {"size": 200}},
]
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
files = await HuggingFaceSource().list_files("u/r")
assert captured["url"] == "https://huggingface.co/api/models/u/r/tree/main"
assert files == [
{"filename": "a/model.safetensors", "size": 300},
{"filename": "b.safetensors", "size": 200},
]
@pytest.mark.asyncio
async def test_huggingface_honours_explicit_revision(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, []
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
await HuggingFaceSource().list_files("u/r", "v2.0")
assert captured["url"].endswith("/tree/v2.0")
@pytest.mark.asyncio
async def test_modelscope_reads_repo_files_api(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, {
"Data": {
"Files": [
# directories are listed too and must be dropped
{"Type": "tree", "Path": "vae", "Size": 0},
{"Type": "blob", "Path": "README.md", "Size": 100},
{"Type": "blob", "Path": "sub/model.safetensors", "Size": 500},
{"Type": "blob", "Path": "model.ckpt", "Size": 200},
]
}
}
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
files = await ModelScopeSource().list_files("u/r")
assert captured["url"] == (
"https://modelscope.cn/api/v1/models/u/r/repo/files?Revision=master"
)
assert files == [
{"filename": "sub/model.safetensors", "size": 500},
{"filename": "model.ckpt", "size": 200},
]
@pytest.mark.asyncio
async def test_missing_repo_is_reported_as_not_found(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 404, None
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
with pytest.raises(ModelSourceError) as excinfo:
await ModelScopeSource().list_files("u/r")
assert excinfo.value.status == 404
assert "not found" in str(excinfo.value)
@pytest.mark.asyncio
async def test_transport_failure_is_reported_as_bad_gateway(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 0, None
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
with pytest.raises(ModelSourceError) as excinfo:
await HuggingFaceSource().list_files("u/r")
assert excinfo.value.status == 502
class TestDownloadUrls:
def test_huggingface_resolve_url(self):
assert HuggingFaceSource().file_download_url("u/r", "sub/f.safetensors") == (
"https://huggingface.co/u/r/resolve/main/sub/f.safetensors"
)
def test_modelscope_resolve_url_defaults_to_master(self):
assert ModelScopeSource().file_download_url("u/r", "sub/f.safetensors") == (
"https://modelscope.cn/models/u/r/resolve/master/sub/f.safetensors"
)
def test_explicit_revision_wins(self):
assert ModelScopeSource().file_download_url("u/r", "f.bin", "v1") == (
"https://modelscope.cn/models/u/r/resolve/v1/f.bin"
)
def test_tensorart_refuses_to_build_a_download_url(self):
source = TensorArtSource()
assert source.supports_download is False
with pytest.raises(ModelSourceError):
source.file_download_url("123", "f.safetensors")
@pytest.mark.asyncio
async def test_tensorart_lists_nothing(self):
assert await TensorArtSource().list_files("123") == []
class TestSourceIdValidation:
@pytest.mark.parametrize(
"source_id",
["u/r", "black-forest-labs/FLUX.1-dev", "AI-ModelScope/stable-diffusion-v1-5"],
)
def test_accepts_repo_ids(self, source_id):
assert is_valid_source_id(source_id) is True
@pytest.mark.parametrize(
"source_id",
[
"",
"noslash",
"a/b/c",
"../etc/passwd",
"u/..",
"u/.",
".hidden/r",
"u/r with space",
"/r",
"u/",
],
)
def test_rejects_unsafe_ids(self, source_id):
assert is_valid_source_id(source_id) is False
class TestDownloadSourceRegistry:
def test_downloadable_sources_excludes_link_only_sites(self):
platforms = {source.platform for source in downloadable_sources()}
assert platforms == {"huggingface", "modelscope"}
def test_get_download_source_rejects_link_only_platform(self):
assert get_download_source("tensorart") is None
assert get_download_source("nope") is None
assert get_download_source("modelscope").platform == "modelscope"
assert get_download_source("huggingface").platform == "huggingface"