Files
ComfyUI-Lora-Manager/tests/frontend/utils/modelSourceUrlDetection.test.js
T
Will Miao 38d4c59b4c feat(download): support ModelScope repositories in the URL downloader
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.
2026-09-14 07:42:51 +08:00

183 lines
6.5 KiB
JavaScript

import { describe, it, expect } from 'vitest';
import { DownloadManager } from '../../../static/js/managers/DownloadManager.js';
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: 'model-source-file',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
revision: 'main',
filename: 'Flux2-Klein-9B-consistency-V2.safetensors',
});
});
it('detects HF resolve URL with subdirectory file', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/user/repo/resolve/main/subdir/model.safetensors'
);
expect(result).toEqual({
type: 'model-source-file',
platform: 'huggingface',
repo: 'user/repo',
revision: 'main',
filename: 'subdir/model.safetensors',
});
});
it('detects HF repo URL (full URL)', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
});
});
it('detects HF repo URL (bare user/repo)', () => {
const result = DownloadManager.detectUrlType('dx8152/Flux2-Klein-9B-Consistency');
expect(result).toEqual({
type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency',
});
});
it('detects HF repo URL with trailing slash', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/user/repo/'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'huggingface',
repo: 'user/repo',
});
});
it('detects HF blob (web preview) URL as resolve', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors'
);
expect(result).toEqual({
type: 'model-source-file',
platform: 'huggingface',
repo: 'Comfy-Org/z_image_turbo',
revision: 'main',
filename: 'split_files/diffusion_models/z_image_turbo_bf16.safetensors',
});
});
it('detects CivitAI URL', () => {
const result = DownloadManager.detectUrlType(
'https://civitai.com/models/123/some-model'
);
expect(result).toEqual({ type: 'civitai' });
});
it('detects CivitAI URL on civitai.red domain', () => {
const result = DownloadManager.detectUrlType(
'https://civitai.red/models/12345/my-model'
);
expect(result).toEqual({ type: 'civitai' });
});
it('detects CivitAI URL on civitai.green domain', () => {
const result = DownloadManager.detectUrlType(
'https://civitai.green/models/67890/another-model'
);
expect(result).toEqual({ type: 'civitai' });
});
it('detects CivArchive URL', () => {
const result = DownloadManager.detectUrlType(
'https://civarchive.com/models/456'
);
expect(result).toEqual({ type: 'civitai' });
});
it('detects direct HTTP URL', () => {
const result = DownloadManager.detectUrlType(
'https://example.com/file.zip'
);
expect(result).toEqual({ type: 'direct-http' });
});
it('returns null for invalid input', () => {
expect(DownloadManager.detectUrlType('')).toBeNull();
expect(DownloadManager.detectUrlType(' ')).toBeNull();
});
it('returns null for unrecognized path', () => {
expect(DownloadManager.detectUrlType('justrandomtext')).toBeNull();
});
it('prefers HF resolve over repo when both match', () => {
const result = DownloadManager.detectUrlType(
'https://huggingface.co/user/repo/resolve/main/file.safetensors'
);
expect(result?.type).toBe('model-source-file');
});
it('prefers CivitAI over HF when both match', () => {
// CivitAI check comes first in detectUrlType
// This URL should be detected as CivitAI, not HF
const result = DownloadManager.detectUrlType(
'https://civitai.com/models/123?huggingface.co/test/repo'
);
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();
});
});