Files
ComfyUI-Lora-Manager/static/js/components/ContextMenu/ModelContextMenuMixin.js
T
Will Miao 5ab0e88abc feat(links): support ModelScope and TensorArt as model sources
A model file could only ever be linked to huggingface.co: `set_hf_url`
validated the URL with a huggingface-only regex, the agent fetched the card
from a hardcoded HF URL, and the readme processor built every relative image
path off `https://huggingface.co/{repo}/resolve/main`. ModelScope publishes the
same model-card convention (README.md + YAML frontmatter, often carrying
`base_model:` and `trigger_words:`) behind a public, key-less API, so the
enrichment pipeline could already serve it - it was the plumbing that was
HF-shaped, not the idea.

Make the external source a first-class, provider-driven concept:

- New `py/services/model_sources/` registry. A `ModelSource` owns URL
  recognition (lenient for stored values, strict for user input), the
  canonical page URL, model-card fetching, the asset base URL and the
  capability flags. `HuggingFaceSource` is the previous logic relocated;
  `ModelScopeSource` reads `/models/{o}/{n}/resolve/{master|main}/README.md`
  and falls back to `/api/v1/models/{o}/{n}/repo`. `TensorArtSource` is
  link-only on purpose: tensor.art answers plain HTTP clients with a
  Cloudflare challenge and its internal API (ap-east-1.tensorart.cloud /
  cn.tensorart.net) rejects every /v1/model/* route with "invalid
  authorization header", so it declares supports_enrichment=False rather than
  failing silently later.
- Metadata gains `source_platform` + `source_url`; `hf_url` stays as a
  read/write alias, written only for Hugging Face, so existing sidecars,
  cached rows and third-party consumers keep working. Normalisation runs at
  the scanner, the persistent cache (both directions, plus two new columns
  behind an ALTER migration) and the linking handler - which is what stops a
  user who switches sources from leaving a stale `hf_url` on a ModelScope
  model.
- The agent pipeline keys off the provider instead of `hf_url`: the fast-fail
  gate now explains *why* a model is skipped (no source / unknown source /
  source without a reachable card), the prompt context exposes
  source_url/source_id/source_label/asset_base_url while still filling the
  legacy hf_url/repo aliases, and the four README image extractors take a
  base_url (defaulting to HF) so relative paths resolve against the right
  site. Version grouping generalises to hf: / ms: / ta: keys.
- `POST /api/lm/set-hf-url` keeps its path and its legacy payload keys but
  accepts `source_url`, validates against every provider and returns the
  platform. `GET /api/lm/model-sources` lets the UI render the supported-site
  list from the server.
- Frontend: a `modelSourceHelpers` mirror of the registry drives the link
  dialog, the card/modal globe (branded "View on ModelScope/TensorArt"), the
  version-group key and the enrichment gate; the versions tab no longer sends
  ms:/ta: keys to the CivitAI API.

TensorArt stays in the list because provenance is worth keeping even when the
card is unreadable - the dialog says so plainly ("Sites that don't expose one
(currently TensorArt) can only be linked") and the context menu disables
enrichment with a matching tooltip, instead of the user getting
"Unsupported URL".

Verified against the real ModelScope API: jj3550945163/Krea-2-LORA returns a
1882-byte card whose frontmatter carries base_model/tags/trigger_words, and
relative images resolve to .../resolve/master/....

Tests: backend 2815 passed; frontend 1130 JS + 91 Vue passed; pytest
tests/i18n and a Jinja compile pass over templates/. The nine locales carry
[TODO: Translate] for the new strings, completed in the next commit.
2026-09-14 07:24:08 +08:00

560 lines
21 KiB
JavaScript

import { showToast, getNSFWLevelName, openExampleImagesFolder } from '../../utils/uiHelpers.js';
import { modalManager } from '../../managers/ModalManager.js';
import { state } from '../../state/index.js';
import { getModelApiClient, resetAndReload } from '../../api/modelApiFactory.js';
import { bulkManager } from '../../managers/BulkManager.js';
import { MODEL_CONFIG } from '../../api/apiConfig.js';
import { translate } from '../../utils/i18nHelpers.js';
import { getNsfwLevelSelector } from '../shared/NsfwLevelSelector.js';
import { classifyModelRelinkUrl } from '../../utils/civitaiUtils.js';
import { parseModelSourceUrl, getModelSourceInfo } from '../../utils/modelSourceHelpers.js';
import { escapeHtml } from '../shared/utils.js';
// Mixin with shared functionality for LoraContextMenu and CheckpointContextMenu
export const ModelContextMenuMixin = {
isExcludedView() {
return state?.pages?.[state.currentPageType]?.viewMode === 'excluded';
},
updateExcludeMenuItem() {
const excludeItem = this.menu?.querySelector('[data-action="exclude"], [data-action="restore"]');
if (!excludeItem) {
return;
}
const isExcludedView = this.isExcludedView();
excludeItem.dataset.action = isExcludedView ? 'restore' : 'exclude';
excludeItem.innerHTML = isExcludedView
? `<i class="fas fa-undo"></i> <span>${translate('loras.contextMenu.restoreModel', {}, 'Restore model')}</span>`
: `<i class="fas fa-eye-slash"></i> <span>${translate('loras.contextMenu.excludeModel', {}, 'Exclude model')}</span>`;
},
async restoreExcludedModel(filePath) {
const restored = await getModelApiClient().unexcludeModel(filePath);
if (!restored) {
return;
}
if (window.pageControls?.exitExcludedView) {
await window.pageControls.exitExcludedView();
} else {
const resetFn = this.resetAndReload || resetAndReload;
if (typeof resetFn === 'function') {
await resetFn(true);
}
}
},
// NSFW Selector methods
initNSFWSelector() {
if (this._nsfwSelectorInitialized) {
return;
}
const selector = getNsfwLevelSelector();
if (!selector) {
console.warn('NSFW selector element not found');
return;
}
this._nsfwSelectorInitialized = true;
this._nsfwSelector = selector;
},
resetNSFWSelectorState() {
// maintained for compatibility; no-op with shared selector
},
showNSFWLevelSelector(x, y, card) {
this.initNSFWSelector();
const selector = this._nsfwSelector || getNsfwLevelSelector();
if (!selector) {
console.warn('NSFW selector not available');
return;
}
// Get current NSFW level
let currentLevel = 0;
try {
const metaData = JSON.parse(card.dataset.meta || '{}');
currentLevel = metaData.preview_nsfw_level || 0;
// Update if we have no recorded level but have a dataset attribute
if (!currentLevel && card.dataset.nsfwLevel) {
currentLevel = parseInt(card.dataset.nsfwLevel, 10) || 0;
}
} catch (err) {
console.error('Error parsing metadata:', err);
}
const filePath = card.dataset.filepath;
selector.show({
currentLevel,
cardPath: filePath,
onSelect: async (level) => {
if (!filePath) return false;
try {
await this.saveModelMetadata(filePath, { preview_nsfw_level: level });
showToast('toast.contextMenu.contentRatingSet', { level: getNSFWLevelName(level) }, 'success');
return true;
} catch (error) {
showToast('toast.contextMenu.contentRatingFailed', { message: error.message }, 'error');
return false;
}
},
onClose: () => this.resetNSFWSelectorState(),
});
},
// Civitai re-linking methods
getModelTypePrefix() {
// Map the mixin model type to its API route prefix; the relink route
// exists for all model types via COMMON_ROUTE_DEFINITIONS.
const prefixMap = {
lora: 'loras',
checkpoint: 'checkpoints',
embedding: 'embeddings',
other: 'other'
};
return prefixMap[this.modelType] || 'loras';
},
showRelinkCivitaiModal() {
const filePath = this.currentCard.dataset.filepath;
if (!filePath) return;
// Set up confirm button handler
const confirmBtn = document.getElementById('confirmRelinkBtn');
const urlInput = document.getElementById('civitaiModelUrl');
const errorDiv = document.getElementById('civitaiModelUrlError');
// Remove previous event listener if exists
if (this._boundRelinkHandler) {
confirmBtn.removeEventListener('click', this._boundRelinkHandler);
}
// Create new bound handler
this._boundRelinkHandler = async () => {
const url = urlInput.value.trim();
const { source, modelId, modelVersionId } = classifyModelRelinkUrl(url);
if (!source || !modelId) {
errorDiv.textContent = 'Invalid URL format. Expected: https://civitai.com/models/{modelId} or https://civarchive.com/models/{modelId}';
return;
}
errorDiv.textContent = '';
modalManager.closeModal('relinkCivitaiModal');
try {
const isCivArchive = source === 'civarchive';
state.loadingManager.showSimpleLoading(
isCivArchive ? 'Re-linking via CivitArchive...' : 'Re-linking to Civitai...'
);
const endpoint = `/api/lm/${this.getModelTypePrefix()}/relink-civitai`;
const payload = {
file_path: filePath,
model_id: modelId,
model_version_id: modelVersionId
};
// Omitted source keeps backend default-provider behaviour; only
// civarchive pins the provider explicitly.
if (isCivArchive) {
payload.source = source;
}
const response = await fetch(endpoint, {
method: 'POST',
headers: {
'Content-Type': 'application/json'
},
body: JSON.stringify(payload)
});
if (!response.ok) {
throw new Error(`Failed to re-link model: ${response.statusText}`);
}
const data = await response.json();
if (data.success) {
showToast(
isCivArchive ? 'toast.contextMenu.linkCivArchSuccess' : 'toast.contextMenu.relinkSuccess',
{},
'success'
);
// Reload the current view to show updated data
await this.resetAndReload();
} else {
throw new Error(data.error || 'Failed to re-link model');
}
} catch (error) {
console.error('Error re-linking model:', error);
showToast('toast.contextMenu.relinkFailed', { message: error.message }, 'error');
} finally {
state.loadingManager.hide();
}
};
// Set new event listener
confirmBtn.addEventListener('click', this._boundRelinkHandler);
// Clear previous input
urlInput.value = '';
errorDiv.textContent = '';
// Show modal
modalManager.showModal('relinkCivitaiModal');
// Auto-focus the URL input field after modal is shown
setTimeout(() => urlInput.focus(), 50);
},
// External model source linking (Hugging Face / ModelScope / TensorArt)
showLinkHfModal() {
const filePath = this.currentCard.dataset.filepath;
if (!filePath) return;
const confirmBtn = document.getElementById('confirmLinkHfBtn');
const urlInput = document.getElementById('hfModelUrl');
const errorDiv = document.getElementById('hfModelUrlError');
if (this._boundLinkHfHandler) {
confirmBtn.removeEventListener('click', this._boundLinkHfHandler);
}
this._boundLinkHfHandler = async () => {
const rawUrl = urlInput.value.trim();
if (!rawUrl) {
errorDiv.textContent = translate(
'modals.linkModelSource.urlRequired',
{},
'Please enter a model page URL.'
);
return;
}
const sourceInfo = parseModelSourceUrl(rawUrl);
if (!sourceInfo) {
errorDiv.textContent = translate(
'modals.linkModelSource.invalidUrl',
{},
'Unsupported URL. Supported sites: Hugging Face, ModelScope, TensorArt.'
);
return;
}
errorDiv.textContent = '';
modalManager.closeModal('linkHfModal');
try {
state.loadingManager.showSimpleLoading(
translate('modals.linkModelSource.linking', {}, 'Linking model source...')
);
const response = await fetch('/api/lm/set-hf-url', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ file_path: filePath, source_url: sourceInfo.url }),
});
if (!response.ok) {
const errData = await response.json().catch(() => ({}));
throw new Error(errData.error || `Request failed: ${response.statusText}`);
}
const data = await response.json();
if (data.success) {
showToast('toast.contextMenu.linkHfSuccess', {}, 'success');
await this.resetAndReload();
} else {
throw new Error(data.error || 'Failed to link model');
}
} catch (error) {
console.error('Error linking model source:', error);
showToast('toast.contextMenu.linkHfFailed', { message: error.message }, 'error');
} finally {
state.loadingManager.hide();
}
};
confirmBtn.addEventListener('click', this._boundLinkHfHandler);
urlInput.value = '';
errorDiv.textContent = '';
modalManager.showModal('linkHfModal');
this._renderSupportedSources();
setTimeout(() => urlInput.focus(), 50);
},
/**
* Refresh the supported-site hints from the server so the dialog reflects
* whatever sources this backend build actually knows about. Falls back to
* the static markup in the template when the request fails.
*/
async _renderSupportedSources() {
const container = document.getElementById('hfSupportedSources');
if (!container) return;
try {
const response = await fetch('/api/lm/model-sources');
if (!response.ok) return;
const sources = await response.json();
if (!Array.isArray(sources) || sources.length === 0) return;
const examples = sources
.map((source) => source?.example_url)
.filter((url) => typeof url === 'string' && url);
if (examples.length === 0) return;
container.innerHTML = examples
.map((url) => `<strong>${escapeHtml(url)}</strong>`)
.join('<br>');
} catch (error) {
console.debug('Failed to load supported model sources:', error);
}
},
// Model metadata enrichment (AI agent) methods
updateEnrichMenuItem(card) {
const enrichItem = this.menu?.querySelector('[data-action="enrich-hf-llm"]');
if (!enrichItem) return;
const model = {
source_url: card.dataset.source_url || '',
source_platform: card.dataset.source_platform || '',
hf_url: card.dataset.hf_url || '',
};
const sourceInfo = getModelSourceInfo(model);
const canEnrich = Boolean(sourceInfo && sourceInfo.supportsEnrichment);
enrichItem.classList.toggle('disabled', !canEnrich);
if (canEnrich) {
enrichItem.title = '';
} else if (!sourceInfo) {
enrichItem.title = translate(
'toast.contextMenu.enrichNeedsSource',
{},
'Link this model to a model source first (Link Model → Link to Model Source)'
);
} else {
enrichItem.title = translate(
'toast.contextMenu.enrichUnsupportedSource',
{ source: sourceInfo.label },
`AI enrichment is not available for ${sourceInfo.label} models`
);
}
},
async enrichWithAgent(filePath) {
const { agentManager } = await import('../../managers/AgentManager.js');
const configured = await agentManager.isLlmConfigured();
if (!configured) {
showToast('toast.agent.llmNotConfigured', {}, 'warning');
return;
}
agentManager.connect();
const progressUI = state.loadingManager.showEnhancedProgress(
'Enriching metadata with AI...'
);
function cleanupCallbacks() {
const pIdx = agentManager.progressCallbacks.indexOf(onProgress);
if (pIdx >= 0) agentManager.progressCallbacks.splice(pIdx, 1);
const cIdx = agentManager.completeCallbacks.indexOf(onComplete);
if (cIdx >= 0) agentManager.completeCallbacks.splice(cIdx, 1);
const eIdx = agentManager.errorCallbacks.indexOf(onError);
if (eIdx >= 0) agentManager.errorCallbacks.splice(eIdx, 1);
}
const onProgress = (data) => {
if (data.status === 'processing' && data.current_path && data.updated_data && Object.keys(data.updated_data).length > 0) {
if (state.virtualScroller?.updateSingleItem) {
state.virtualScroller.updateSingleItem(data.current_path, data.updated_data);
}
const pct = data.total > 0 ? Math.floor((data.processed / data.total) * 100) : 0;
const name = data.current_path.split('/').pop();
progressUI.updateProgress(pct, name, `Processing ${name}`);
}
};
agentManager.onProgress(onProgress);
const onComplete = (data) => {
cleanupCallbacks();
if (data.status === 'completed') {
progressUI.complete(data.summary || 'Enrich complete');
showToast('toast.agent.enrichComplete', { summary: data.summary || 'Done' }, 'success');
}
};
agentManager.onComplete(onComplete);
const onError = (data) => {
cleanupCallbacks();
state.loadingManager.hide();
showToast('toast.agent.enrichFailed', { error: data.error || 'Unknown error' }, 'error');
};
agentManager.onError(onError);
try {
await agentManager.executeSkill('enrich_hf_metadata', [filePath]);
} catch (error) {
cleanupCallbacks();
state.loadingManager.hide();
showToast('toast.agent.enrichFailed', { error: error.message }, 'error');
}
},
parseModelId(value) {
if (value === undefined || value === null || value === '') {
return null;
}
const parsed = Number.parseInt(value, 10);
return Number.isNaN(parsed) ? null : parsed;
},
getModelIdFromCard(card) {
if (!card) {
return null;
}
if (card.dataset?.meta) {
try {
const meta = JSON.parse(card.dataset.meta);
const metaValue = this.parseModelId(meta?.modelId);
if (metaValue !== null) {
return metaValue;
}
} catch (error) {
console.warn('Unable to parse card metadata for model ID', error);
}
}
return null;
},
async checkUpdatesForCurrentModel() {
const card = this.currentCard;
if (!card) {
return;
}
const modelId = this.getModelIdFromCard(card);
const typeConfig = MODEL_CONFIG[this.modelType] || {};
const typeLabel = (typeConfig.displayName || 'Model').toLowerCase();
if (modelId === null) {
showToast('toast.models.bulkUpdatesMissing', { type: typeLabel }, 'warning');
return;
}
const apiClient = getModelApiClient();
const loadingMessage = translate(
'toast.models.bulkUpdatesChecking',
{ count: 1, type: typeLabel },
`Checking selected ${typeLabel}(s) for updates...`
);
state.loadingManager.showSimpleLoading(loadingMessage);
try {
const response = await apiClient.refreshUpdatesForModels([modelId]);
const records = Array.isArray(response?.records) ? response.records : [];
const updatesCount = records.length;
if (updatesCount > 0) {
showToast('toast.models.bulkUpdatesSuccess', { count: updatesCount, type: typeLabel }, 'success');
} else {
showToast('toast.models.bulkUpdatesNone', { type: typeLabel }, 'info');
}
const resetFn = this.resetAndReload || resetAndReload;
if (typeof resetFn === 'function') {
await resetFn(false);
}
} catch (error) {
console.error('Error checking updates for model:', error);
showToast(
'toast.models.bulkUpdatesFailed',
{ type: typeLabel, message: error?.message ?? 'Unknown error' },
'error'
);
} finally {
state.loadingManager.hide();
state.loadingManager.restoreProgressBar();
}
},
// Common action handlers
handleCommonMenuActions(action) {
switch(action) {
case 'preview':
openExampleImagesFolder(this.currentCard.dataset.sha256);
return true;
case 'download-examples':
this.downloadExampleImages(false);
return true;
case 'download-examples-force':
this.downloadExampleImages(true);
return true;
case 'civitai':
// 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 {
showToast('toast.contextMenu.fetchMetadataFirst', {}, 'info');
}
} else {
showToast('toast.contextMenu.noCivitaiInfo', {}, 'info');
}
return true;
case 'relink-civitai':
this.showRelinkCivitaiModal();
return true;
case 'link-hf':
this.showLinkHfModal();
return true;
case 'enrich-hf-llm':
this.enrichWithAgent(this.currentCard.dataset.filepath);
return true;
case 'set-nsfw':
this.showNSFWLevelSelector(null, null, this.currentCard);
return true;
case 'check-updates':
this.checkUpdatesForCurrentModel();
return true;
default:
return false;
}
},
// Download example images method
async downloadExampleImages(force = false) {
const modelHash = this.currentCard.dataset.sha256;
if (!modelHash) {
showToast('toast.contextMenu.missingHash', {}, 'error');
return;
}
try {
const apiClient = getModelApiClient();
await apiClient.downloadExampleImages([modelHash], null, { force });
} catch (error) {
console.error('Error downloading example images:', error);
}
}
};