mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
fix(links): let CivitAI and HuggingFace links coexist (#1094)
A model could have CivitAI metadata and a HuggingFace link at the same time, but only one of the two "View on ..." entries ever rendered, because both the model modal and the card globe asked the `from_civitai` provenance flag which source to show. `set_hf_url` wrote `false` and a CivitAI refresh wrote `true`, so whichever ran last erased the other: linking HF hid "View on CivitAI" even though the civitai payload was still in the sidecar, and (on the card) a later refresh pointed the single globe icon back at CivitAI, hiding the HF entry. Decide the links from the data itself instead: - `set_hf_url` no longer touches `from_civitai`; it records where the metadata came from, and HF provenance is already tracked by `hf_url`. - Add `hasCivitaiSource(civitai)` in the shared card/modal utils and gate the modal's CivitAI link, the card globe (title, enabled state, click target, new `data-has_civitai`) and the context-menu `civitai` action on actual CivitAI data (`modelId` / `model_id` / `id`). A dual-source model now shows both links, and a CivitAI-only model with no `hf_url` stays as before. - Agent HF enrichment (`PostProcessor.is_hf_model`) keyed off `not from_civitai`, which stopped being a synonym for "has an HF source" once both sources can coexist (and already broke after a CivitAI refresh flipped the flag back to true). Key it off `hf_url` directly; the post-processor tests move to that discriminator and gain a dual-source case. Regression tests: the set-hf-url handler preserves civitai + `from_civitai` and no longer forces the flag false, the modal renders both links (including with `from_civitai: false`), and the card globe targets/opens the right source and is disabled when neither is available. Backend: 2749 passed. Frontend: 1098 JS + 91 Vue tests passed.
This commit is contained in:
Will Miao
@@ -0,0 +1,194 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
|
||||
const {
|
||||
MODEL_CARD_MODULE,
|
||||
STATE_MODULE,
|
||||
UI_HELPERS_MODULE,
|
||||
I18N_MODULE,
|
||||
API_CONFIG_MODULE,
|
||||
API_FACTORY_MODULE,
|
||||
} = vi.hoisted(() => ({
|
||||
MODEL_CARD_MODULE: new URL('../../../static/js/components/shared/ModelCard.js', import.meta.url).pathname,
|
||||
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||
I18N_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||
API_CONFIG_MODULE: new URL('../../../static/js/api/apiConfig.js', import.meta.url).pathname,
|
||||
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||
}));
|
||||
|
||||
vi.mock(STATE_MODULE, () => ({
|
||||
state: {
|
||||
settings: {
|
||||
blur_mature_content: false,
|
||||
model_name_display: 'model_name',
|
||||
},
|
||||
global: {
|
||||
settings: {
|
||||
model_name_display: 'model_name',
|
||||
group_by_model: false,
|
||||
display_density: 'default',
|
||||
model_card_footer_action: 'example_images',
|
||||
},
|
||||
},
|
||||
pages: {
|
||||
other: {
|
||||
previewVersions: new Map(),
|
||||
sortBy: 'name',
|
||||
},
|
||||
},
|
||||
bulkMode: false,
|
||||
selectedModels: new Set(),
|
||||
selectedLoras: new Set(),
|
||||
},
|
||||
getCurrentPageState: vi.fn(() => ({
|
||||
sortBy: 'name',
|
||||
previewVersions: new Map(),
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||
showToast: vi.fn(),
|
||||
openCivitai: vi.fn(),
|
||||
openHuggingFace: vi.fn(),
|
||||
copyToClipboard: vi.fn(),
|
||||
copyLoraSyntax: vi.fn(),
|
||||
sendLoraToWorkflow: vi.fn(),
|
||||
sendEmbeddingToWorkflow: vi.fn(),
|
||||
openExampleImagesFolder: vi.fn(),
|
||||
buildLoraSyntax: vi.fn(),
|
||||
sendModelPathToWorkflow: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(I18N_MODULE, () => ({
|
||||
translate: vi.fn((key, params, fallback) => (typeof fallback === 'string' ? fallback : key)),
|
||||
}));
|
||||
|
||||
vi.mock(API_CONFIG_MODULE, () => ({
|
||||
MODEL_TYPES: { LORA: 'loras', CHECKPOINT: 'checkpoints', EMBEDDING: 'embeddings', OTHER: 'other' },
|
||||
}));
|
||||
|
||||
vi.mock(API_FACTORY_MODULE, () => ({
|
||||
getModelApiClient: vi.fn(() => ({})),
|
||||
}));
|
||||
|
||||
function makeModel(overrides = {}) {
|
||||
return {
|
||||
sha256: 'abc123',
|
||||
file_path: '/models/loras/linked.safetensors',
|
||||
model_name: 'Linked LoRA',
|
||||
file_name: 'linked',
|
||||
folder: '',
|
||||
modified: 1234567890,
|
||||
file_size: 1024,
|
||||
notes: '',
|
||||
base_model: '',
|
||||
favorite: false,
|
||||
exclude: false,
|
||||
from_civitai: true,
|
||||
hf_url: '',
|
||||
update_available: false,
|
||||
skip_metadata_refresh: false,
|
||||
preview_url: '',
|
||||
preview_nsfw_level: 0,
|
||||
tags: [],
|
||||
civitai: {},
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function mountCard(createModelCard, model) {
|
||||
document.body.innerHTML = '<div id="modelGrid"></div>';
|
||||
const card = createModelCard(model, 'loras');
|
||||
document.getElementById('modelGrid').appendChild(card);
|
||||
return card;
|
||||
}
|
||||
|
||||
describe('ModelCard source globe (#1094)', () => {
|
||||
let createModelCard;
|
||||
let setupModelCardEventDelegation;
|
||||
let openCivitai;
|
||||
let openHuggingFace;
|
||||
|
||||
beforeEach(async () => {
|
||||
document.body.innerHTML = '';
|
||||
({ createModelCard, setupModelCardEventDelegation } = await import(MODEL_CARD_MODULE));
|
||||
({ openCivitai, openHuggingFace } = await import(UI_HELPERS_MODULE));
|
||||
openCivitai.mockReset();
|
||||
openHuggingFace.mockReset();
|
||||
});
|
||||
|
||||
it('points the globe at CivitAI when CivitAI data is present alongside an HF link', () => {
|
||||
const card = mountCard(
|
||||
createModelCard,
|
||||
makeModel({
|
||||
civitai: { id: 111, modelId: 222, name: 'v1' },
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
})
|
||||
);
|
||||
|
||||
expect(card.dataset.has_civitai).toBe('true');
|
||||
expect(card.querySelector('.fa-globe').getAttribute('title')).toBe('View on Civitai');
|
||||
});
|
||||
|
||||
it('keeps the CivitAI globe target when from_civitai is false but CivitAI data exists', () => {
|
||||
// Regression for case 1: linking HF no longer hides CivitAI.
|
||||
const card = mountCard(
|
||||
createModelCard,
|
||||
makeModel({
|
||||
from_civitai: false,
|
||||
civitai: { id: 111, modelId: 222 },
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
})
|
||||
);
|
||||
|
||||
expect(card.dataset.has_civitai).toBe('true');
|
||||
expect(card.querySelector('.fa-globe').getAttribute('title')).toBe('View on Civitai');
|
||||
});
|
||||
|
||||
it('points the globe at HuggingFace for an HF-only model', () => {
|
||||
const card = mountCard(
|
||||
createModelCard,
|
||||
makeModel({ from_civitai: false, civitai: {}, hf_url: 'https://huggingface.co/user/repo' })
|
||||
);
|
||||
|
||||
expect(card.dataset.has_civitai).toBe('false');
|
||||
expect(card.querySelector('.fa-globe').getAttribute('title')).toBe('View on Hugging Face');
|
||||
});
|
||||
|
||||
it('disables the globe when there is no CivitAI data and no HF link', () => {
|
||||
const card = mountCard(createModelCard, makeModel({ civitai: {} }));
|
||||
const globe = card.querySelector('.fa-globe');
|
||||
|
||||
expect(card.dataset.has_civitai).toBe('false');
|
||||
expect(globe.getAttribute('style')).toContain('cursor: not-allowed');
|
||||
});
|
||||
|
||||
it('opens CivitAI when the globe is clicked on a dual-source model', () => {
|
||||
const model = makeModel({
|
||||
civitai: { id: 111, modelId: 222 },
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
});
|
||||
const card = mountCard(createModelCard, model);
|
||||
setupModelCardEventDelegation('loras');
|
||||
|
||||
card.querySelector('.fa-globe').dispatchEvent(new MouseEvent('click', { bubbles: true }));
|
||||
|
||||
expect(openCivitai).toHaveBeenCalledWith(model.file_path);
|
||||
expect(openHuggingFace).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('opens HuggingFace when the globe is clicked on an HF-only model', () => {
|
||||
const model = makeModel({
|
||||
from_civitai: false,
|
||||
civitai: {},
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
});
|
||||
const card = mountCard(createModelCard, model);
|
||||
setupModelCardEventDelegation('loras');
|
||||
|
||||
card.querySelector('.fa-globe').dispatchEvent(new MouseEvent('click', { bubbles: true }));
|
||||
|
||||
expect(openHuggingFace).toHaveBeenCalledWith('https://huggingface.co/user/repo');
|
||||
expect(openCivitai).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,191 @@
|
||||
import { describe, it, beforeEach, expect, vi } from 'vitest';
|
||||
|
||||
const {
|
||||
MODAL_MODULE,
|
||||
API_FACTORY,
|
||||
UI_HELPERS_MODULE,
|
||||
MODAL_MANAGER_MODULE,
|
||||
SHOWCASE_MODULE,
|
||||
MODEL_TAGS_MODULE,
|
||||
UTILS_MODULE,
|
||||
TRIGGER_WORDS_MODULE,
|
||||
PRESET_TAGS_MODULE,
|
||||
MODEL_VERSIONS_MODULE,
|
||||
RECIPE_TAB_MODULE,
|
||||
I18N_HELPERS_MODULE,
|
||||
} = vi.hoisted(() => ({
|
||||
MODAL_MODULE: new URL('../../../static/js/components/shared/ModelModal.js', import.meta.url).pathname,
|
||||
API_FACTORY: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
|
||||
SHOWCASE_MODULE: new URL('../../../static/js/components/shared/showcase/ShowcaseView.js', import.meta.url).pathname,
|
||||
MODEL_TAGS_MODULE: new URL('../../../static/js/components/shared/ModelTags.js', import.meta.url).pathname,
|
||||
UTILS_MODULE: new URL('../../../static/js/components/shared/utils.js', import.meta.url).pathname,
|
||||
TRIGGER_WORDS_MODULE: new URL('../../../static/js/components/shared/TriggerWords.js', import.meta.url).pathname,
|
||||
PRESET_TAGS_MODULE: new URL('../../../static/js/components/shared/PresetTags.js', import.meta.url).pathname,
|
||||
MODEL_VERSIONS_MODULE: new URL('../../../static/js/components/shared/ModelVersionsTab.js', import.meta.url).pathname,
|
||||
RECIPE_TAB_MODULE: new URL('../../../static/js/components/shared/RecipeTab.js', import.meta.url).pathname,
|
||||
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||
}));
|
||||
|
||||
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||
showToast: vi.fn(),
|
||||
openCivitai: vi.fn(),
|
||||
copyToClipboard: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(MODAL_MANAGER_MODULE, () => ({
|
||||
modalManager: {
|
||||
showModal: vi.fn((id, html) => {
|
||||
document.body.innerHTML = `<div id="${id}">${html}</div>`;
|
||||
}),
|
||||
closeModal: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock(SHOWCASE_MODULE, () => ({
|
||||
scrollToTop: vi.fn(),
|
||||
loadExampleImages: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(MODEL_TAGS_MODULE, () => ({
|
||||
setupTagEditMode: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(UTILS_MODULE, async (importOriginal) => {
|
||||
const actual = await importOriginal();
|
||||
return {
|
||||
...actual,
|
||||
renderCompactTags: vi.fn(() => ''),
|
||||
setupTagTooltip: vi.fn(),
|
||||
formatFileSize: vi.fn(() => '1 MB'),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock(TRIGGER_WORDS_MODULE, () => ({
|
||||
renderTriggerWords: vi.fn(() => ''),
|
||||
setupTriggerWordsEditMode: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(PRESET_TAGS_MODULE, () => ({
|
||||
parsePresets: vi.fn(() => ({})),
|
||||
renderPresetTags: vi.fn(() => ''),
|
||||
}));
|
||||
|
||||
vi.mock(MODEL_VERSIONS_MODULE, () => ({
|
||||
initVersionsTab: vi.fn(() => ({
|
||||
load: vi.fn().mockResolvedValue(undefined),
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock(RECIPE_TAB_MODULE, () => ({
|
||||
loadRecipesForModel: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(I18N_HELPERS_MODULE, () => ({
|
||||
translate: vi.fn((_, __, fallback) => fallback || ''),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/apiConfig.js', () => ({
|
||||
MODEL_TYPES: {
|
||||
LORA: 'loras',
|
||||
CHECKPOINT: 'checkpoints',
|
||||
EMBEDDING: 'embeddings',
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock(API_FACTORY, () => ({
|
||||
getModelApiClient: vi.fn(),
|
||||
}));
|
||||
|
||||
function makeModel(overrides = {}) {
|
||||
return {
|
||||
model_name: 'Linked Model',
|
||||
file_path: 'models/linked.safetensors',
|
||||
file_name: 'linked.safetensors',
|
||||
sha256: 'a'.repeat(64),
|
||||
from_civitai: true,
|
||||
civitai: {},
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe('Model modal source links (#1094)', () => {
|
||||
beforeEach(async () => {
|
||||
document.body.innerHTML = '';
|
||||
const { getModelApiClient } = await import(API_FACTORY);
|
||||
getModelApiClient.mockReset();
|
||||
getModelApiClient.mockReturnValue({
|
||||
fetchModelMetadata: vi.fn().mockResolvedValue(null),
|
||||
saveModelMetadata: vi.fn(),
|
||||
});
|
||||
});
|
||||
|
||||
async function renderModal(model) {
|
||||
const { showModelModal } = await import(MODAL_MODULE);
|
||||
await showModelModal(model, 'loras');
|
||||
}
|
||||
|
||||
const civitaiLink = () => document.querySelector('[data-action="view-civitai"]');
|
||||
const hfLink = () => document.querySelector('[data-action="view-huggingface"]');
|
||||
|
||||
it('renders both links when the model has CivitAI data and an HF link', async () => {
|
||||
await renderModal(
|
||||
makeModel({
|
||||
civitai: { id: 111, modelId: 222, name: 'v1' },
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
})
|
||||
);
|
||||
|
||||
expect(civitaiLink()).not.toBeNull();
|
||||
expect(hfLink()).not.toBeNull();
|
||||
expect(hfLink().dataset.hfUrl).toBe('https://huggingface.co/user/repo');
|
||||
});
|
||||
|
||||
it('keeps the CivitAI link after linking HF even when from_civitai is false', async () => {
|
||||
// Regression for case 1: set_hf_url used to flip from_civitai to false,
|
||||
// which hid the CivitAI link despite the model still having CivitAI data.
|
||||
await renderModal(
|
||||
makeModel({
|
||||
from_civitai: false,
|
||||
civitai: { id: 111, modelId: 222, name: 'v1' },
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
})
|
||||
);
|
||||
|
||||
expect(civitaiLink()).not.toBeNull();
|
||||
expect(hfLink()).not.toBeNull();
|
||||
});
|
||||
|
||||
it('renders only the HF link for an HF-only model', async () => {
|
||||
await renderModal(
|
||||
makeModel({
|
||||
from_civitai: false,
|
||||
civitai: {},
|
||||
hf_url: 'https://huggingface.co/user/repo',
|
||||
})
|
||||
);
|
||||
|
||||
expect(civitaiLink()).toBeNull();
|
||||
expect(hfLink()).not.toBeNull();
|
||||
});
|
||||
|
||||
it('renders the CivitAI link from civitai.model_id when modelId is absent', async () => {
|
||||
await renderModal(
|
||||
makeModel({
|
||||
from_civitai: false,
|
||||
civitai: { id: 111, model_id: 222 },
|
||||
})
|
||||
);
|
||||
|
||||
expect(civitaiLink()).not.toBeNull();
|
||||
expect(hfLink()).toBeNull();
|
||||
});
|
||||
|
||||
it('renders neither link when there is no CivitAI data and no HF link', async () => {
|
||||
await renderModal(makeModel({ civitai: {} }));
|
||||
|
||||
expect(civitaiLink()).toBeNull();
|
||||
expect(hfLink()).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,156 @@
|
||||
"""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()
|
||||
@@ -127,7 +127,11 @@ class TestEnrichHfMetadata:
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=llm,
|
||||
metadata={"base_model": "SD 1.5", "from_civitai": False},
|
||||
metadata={
|
||||
"base_model": "SD 1.5",
|
||||
"from_civitai": False,
|
||||
"hf_url": "https://huggingface.co/user/repo",
|
||||
},
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert applied["base_model"] == "Flux.1 D"
|
||||
@@ -184,14 +188,17 @@ class TestEnrichHfMetadata:
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=llm,
|
||||
metadata={"from_civitai": False},
|
||||
metadata={
|
||||
"from_civitai": False,
|
||||
"hf_url": "https://huggingface.co/user/repo",
|
||||
},
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert applied["civitai"]["description"] == "A short summary"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_description_skipped_for_civitai_model(self, processor):
|
||||
"""short_description NOT written for CivitAI models (has own description)."""
|
||||
async def test_short_description_skipped_without_hf_url(self, processor):
|
||||
"""short_description NOT written when the model has no HF source."""
|
||||
llm = {**self.MIN_LLM_OUTPUT, "short_description": "A short summary"}
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
@@ -221,7 +228,10 @@ class TestEnrichHfMetadata:
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=self.MIN_LLM_OUTPUT,
|
||||
metadata={"from_civitai": False},
|
||||
metadata={
|
||||
"from_civitai": False,
|
||||
"hf_url": "https://huggingface.co/user/repo",
|
||||
},
|
||||
readme_content="# Hello\n\nThis is **bold**.",
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
@@ -229,8 +239,8 @@ class TestEnrichHfMetadata:
|
||||
assert "<strong>bold</strong>" in applied.get("modelDescription", "")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_readme_content_skipped_for_civitai_model(self, processor):
|
||||
"""README content NOT converted for CivitAI models."""
|
||||
async def test_readme_content_skipped_without_hf_url(self, processor):
|
||||
"""README content NOT converted when the model has no HF source."""
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview", return_value=None),
|
||||
@@ -283,8 +293,31 @@ Content
|
||||
assert images[0]["meta"]["prompt"] == "a cat"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gallery_images_skipped_for_civitai_model(self, processor):
|
||||
"""Gallery images NOT extracted for CivitAI models."""
|
||||
async def test_gallery_images_skipped_without_hf_url(self, processor):
|
||||
"""Gallery images NOT extracted when the model has no HF source."""
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview", return_value=None),
|
||||
mock.patch("py.metadata_ops.refresh_cache"),
|
||||
):
|
||||
await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=self.MIN_LLM_OUTPUT,
|
||||
metadata={"from_civitai": True},
|
||||
readme_content="---\nwidget:\n- text: a\n output:\n url: x.png\n---\n",
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
civitai = applied.get("civitai", {})
|
||||
assert "images" not in civitai
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gallery_images_extracted_for_civitai_linked_model(self, processor):
|
||||
"""A model may be on CivitAI and HuggingFace at once (#1094).
|
||||
|
||||
HF enrichment is gated on ``hf_url``, not on ``from_civitai``, so the
|
||||
README gallery is still applied when both sources are present.
|
||||
"""
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview", return_value=None),
|
||||
@@ -301,8 +334,11 @@ Content
|
||||
readme_content="---\nwidget:\n- text: a\n output:\n url: x.png\n---\n",
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
civitai = applied.get("civitai", {})
|
||||
assert "images" not in civitai
|
||||
images = applied.get("civitai", {}).get("images", [])
|
||||
assert len(images) == 1
|
||||
assert images[0]["url"] == (
|
||||
"https://huggingface.co/user/repo/resolve/main/x.png"
|
||||
)
|
||||
|
||||
# -- tags ------------------------------------------------------------
|
||||
|
||||
|
||||
Reference in New Issue
Block a user