diff --git a/py/routes/handlers/hf_handlers.py b/py/routes/handlers/hf_handlers.py index 24a3f23e..dd121361 100644 --- a/py/routes/handlers/hf_handlers.py +++ b/py/routes/handlers/hf_handlers.py @@ -240,7 +240,10 @@ class HfHandler: }) existing["hf_url"] = hf_url - existing["from_civitai"] = False + # NOTE: deliberately do NOT touch `from_civitai` here. It records + # where the metadata came from, and the UI must show the CivitAI + # link whenever CivitAI data is present — linking HuggingFace must + # not hide it (#1094). HF provenance is tracked via `hf_url`. await MetadataManager.save_metadata(file_path, existing) await _add_to_scanner_cache(file_path, existing) diff --git a/py/services/agent/post_processor.py b/py/services/agent/post_processor.py index c863c5ff..37a82a48 100644 --- a/py/services/agent/post_processor.py +++ b/py/services/agent/post_processor.py @@ -92,7 +92,10 @@ class PostProcessor: preview_downloaded = False # -- Determine whether this is an HF-sourced model ----------------- - is_hf_model = not metadata.get("from_civitai", True) + # Key off `hf_url` directly: `from_civitai` records provenance and can + # be true for a model that is also linked to HuggingFace (both sources + # coexist, see #1094), so it must not gate HF enrichment. + is_hf_model = bool(metadata.get("hf_url", "")) # -- Collect updates ----------------------------------------------- updates: Dict[str, Any] = {} diff --git a/static/js/components/ContextMenu/ModelContextMenuMixin.js b/static/js/components/ContextMenu/ModelContextMenuMixin.js index 4fcf96c8..5e163e7d 100644 --- a/static/js/components/ContextMenu/ModelContextMenuMixin.js +++ b/static/js/components/ContextMenu/ModelContextMenuMixin.js @@ -446,7 +446,10 @@ export const ModelContextMenuMixin = { this.downloadExampleImages(true); return true; case 'civitai': - if (this.currentCard.dataset.from_civitai === 'true') { + // Gate on actual CivitAI data (not the `from_civitai` flag) so + // that linking HuggingFace does not make the model look like it + // has no CivitAI info (#1094). + if (this.currentCard.dataset.has_civitai === 'true') { if (this.currentCard.querySelector('.fa-globe')) { this.currentCard.querySelector('.fa-globe').click(); } else { diff --git a/static/js/components/shared/ModelCard.js b/static/js/components/shared/ModelCard.js index 67dcee31..3c5cffe7 100644 --- a/static/js/components/shared/ModelCard.js +++ b/static/js/components/shared/ModelCard.js @@ -1,6 +1,7 @@ import { showToast, openCivitai, openHuggingFace, copyToClipboard, copyLoraSyntax, sendLoraToWorkflow, sendEmbeddingToWorkflow, openExampleImagesFolder, buildLoraSyntax, sendModelPathToWorkflow } from '../../utils/uiHelpers.js'; import { state, getCurrentPageState } from '../../state/index.js'; import { showModelModal } from './ModelModal.js'; +import { hasCivitaiSource } from './utils.js'; import { bulkManager } from '../../managers/BulkManager.js'; import { modalManager } from '../../managers/ModalManager.js'; import { NSFW_LEVELS, getBaseModelAbbreviation, getSubTypeAbbreviation, getMatureBlurThreshold, MODEL_SUBTYPE_DISPLAY_NAMES, MODEL_CARD_DRAG_MIME_TYPE } from '../../utils/constants.js'; @@ -63,7 +64,10 @@ function handleModelCardEvent_internal(event, modelType) { if (event.target.closest('.fa-globe')) { event.stopPropagation(); - if (card.dataset.from_civitai === 'true') { + // CivitAI wins when the model actually has CivitAI data; otherwise fall + // back to HuggingFace. Relying on `from_civitai` here made the two + // sources mutually exclusive whenever one of them was (re)linked (#1094). + if (card.dataset.has_civitai === 'true') { openCivitai(card.dataset.filepath); } else if (card.dataset.hf_url) { openHuggingFace(card.dataset.hf_url); @@ -478,6 +482,9 @@ export function createModelCard(model, modelType) { card.dataset.modified = model.modified; card.dataset.file_size = model.file_size; card.dataset.from_civitai = model.from_civitai; + // Independent of `from_civitai`: a model can have both CivitAI data and an + // HF link, and the card globe must keep pointing at CivitAI when it does. + card.dataset.has_civitai = hasCivitaiSource(model.civitai) ? 'true' : 'false'; card.dataset.usage_count = String(model.usage_count); card.dataset.notes = model.notes || ''; card.dataset.base_model = model.base_model || 'Unknown'; @@ -600,12 +607,13 @@ export function createModelCard(model, modelType) { const favoriteTitle = isFavorite ? translate('modelCard.actions.removeFromFavorites', {}, 'Remove from favorites') : translate('modelCard.actions.addToFavorites', {}, 'Add to favorites'); - const globeTitle = model.from_civitai ? + const hasCivitai = hasCivitaiSource(model.civitai); + const globeTitle = hasCivitai ? translate('modelCard.actions.viewOnCivitai', {}, 'View on Civitai') : model.hf_url ? translate('modelCard.actions.viewOnHuggingFace', {}, 'View on Hugging Face') : translate('modelCard.actions.notAvailableFromCivitai', {}, 'Not available from Civitai'); - const globeEnabled = model.from_civitai || !!model.hf_url; + const globeEnabled = hasCivitai || !!model.hf_url; let sendTitle; let copyTitle; if (modelType === MODEL_TYPES.LORA) { diff --git a/static/js/components/shared/ModelModal.js b/static/js/components/shared/ModelModal.js index 81893625..4177e24b 100644 --- a/static/js/components/shared/ModelModal.js +++ b/static/js/components/shared/ModelModal.js @@ -14,7 +14,7 @@ import { } from './ModelMetadata.js'; import { setupTagEditMode } from './ModelTags.js'; import { getModelApiClient } from '../../api/modelApiFactory.js'; -import { renderCompactTags, setupTagTooltip, formatFileSize, escapeAttribute, escapeHtml } from './utils.js'; +import { renderCompactTags, setupTagTooltip, formatFileSize, escapeAttribute, escapeHtml, hasCivitaiSource } from './utils.js'; import { renderTriggerWords, setupTriggerWordsEditMode } from './TriggerWords.js'; import { parsePresets, renderPresetTags } from './PresetTags.js'; import { initVersionsTab } from './ModelVersionsTab.js'; @@ -389,7 +389,11 @@ export async function showModelModal(model, modelType) { const licenseIcons = useNewIcons ? renderNewLicenseIcons(modelWithFullData) : renderLicenseIcons(modelWithFullData); - const viewOnCivitaiAction = modelWithFullData.from_civitai ? ` + // Gate the CivitAI link on actual CivitAI data, not the `from_civitai` + // provenance flag: a model can be linked to HuggingFace and to CivitAI at + // the same time, and both links must coexist (#1094). + const hasCivitai = hasCivitaiSource(modelWithFullData.civitai); + const viewOnCivitaiAction = hasCivitai ? `
${translate('modals.model.actions.viewOnCivitaiText', {}, 'View on Civitai')}
`.trim() : ''; diff --git a/static/js/components/shared/utils.js b/static/js/components/shared/utils.js index e8b4a38b..42fd24e5 100644 --- a/static/js/components/shared/utils.js +++ b/static/js/components/shared/utils.js @@ -36,6 +36,24 @@ export function formatFileSize(bytes) { return `${size.toFixed(1)} ${units[unitIndex]}`; } +/** + * Whether a model has usable CivitAI metadata to link to. + * + * CivitAI links must be gated on the presence of actual CivitAI data rather + * than the `from_civitai` provenance flag: linking a model to HuggingFace used + * to flip `from_civitai` to false, which hid the CivitAI link even though the + * model still had CivitAI metadata. See issue #1094. + * + * @param {Object} [civitaiData] - The model's `civitai` payload + * @returns {boolean} True when a CivitAI model/version id is available + */ +export function hasCivitaiSource(civitaiData) { + if (!civitaiData || typeof civitaiData !== 'object') return false; + return Boolean( + civitaiData.modelId ?? civitaiData.model_id ?? civitaiData.id + ); +} + /** * Render compact tags * @param {Array} tags - Array of tags diff --git a/tests/frontend/components/modelCard.sourceGlobe.test.js b/tests/frontend/components/modelCard.sourceGlobe.test.js new file mode 100644 index 00000000..37b80b48 --- /dev/null +++ b/tests/frontend/components/modelCard.sourceGlobe.test.js @@ -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 = '
'; + 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(); + }); +}); diff --git a/tests/frontend/components/modelModal.sourceLinks.test.js b/tests/frontend/components/modelModal.sourceLinks.test.js new file mode 100644 index 00000000..9eef1dc0 --- /dev/null +++ b/tests/frontend/components/modelModal.sourceLinks.test.js @@ -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 = `
${html}
`; + }), + 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(); + }); +}); diff --git a/tests/routes/test_hf_handlers.py b/tests/routes/test_hf_handlers.py new file mode 100644 index 00000000..683e714b --- /dev/null +++ b/tests/routes/test_hf_handlers.py @@ -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() diff --git a/tests/services/test_post_processor.py b/tests/services/test_post_processor.py index 8da7a431..09625a2a 100644 --- a/tests/services/test_post_processor.py +++ b/tests/services/test_post_processor.py @@ -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 "bold" 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 ------------------------------------------------------------