feat: add HuggingFace mirror URL setting for missing-model downloads

Author: xingchen-mu

src/platform/missingModel/missingModelDownload.test.ts

@@ -4,26 +4,43 @@ import {
   clearMetadataCache,
   downloadModel,
   fetchModelMetadata,
+  HUGGINGFACE_MIRROR_SETTING_ID,
   isModelDownloadable,
   isTrustedHuggingFaceUrl,
   openGatedRepoPage,
+  resolveHuggingFaceUrl,
   toBrowsableUrl
 } from './missingModelDownload'
 
-const { fetchMock, mockIsDesktop, mockSidebarTabStore, mockStartDownload } =
-  vi.hoisted(() => ({
-    fetchMock: vi.fn(),
-    mockIsDesktop: { value: false },
-    mockSidebarTabStore: { activeSidebarTabId: null as string | null },
-    mockStartDownload: vi.fn()
-  }))
+const {
+  fetchMock,
+  mockHuggingFaceMirror,
+  mockIsDesktop,
+  mockSidebarTabStore,
+  mockStartDownload
+} = vi.hoisted(() => ({
+  fetchMock: vi.fn(),
+  mockHuggingFaceMirror: { value: '' as string | undefined },
+  mockIsDesktop: { value: false },
+  mockSidebarTabStore: { activeSidebarTabId: null as string | null },
+  mockStartDownload: vi.fn()
+}))
 
 vi.mock('@/platform/distribution/types', () => ({
   get isDesktop() {
     return mockIsDesktop.value
   }
 }))
 
+vi.mock('@/platform/settings/settingStore', () => ({
+  useSettingStore: () => ({
+    get: (key: string) =>
+      key === HUGGINGFACE_MIRROR_SETTING_ID
+        ? mockHuggingFaceMirror.value
+        : undefined
+  })
+}))
+
 vi.mock('@/stores/electronDownloadStore', () => ({
   useElectronDownloadStore: () => ({
     start: mockStartDownload
@@ -38,6 +55,7 @@ beforeEach(() => {
   vi.stubGlobal('fetch', fetchMock)
   clearMetadataCache()
   delete window.__comfyDesktop2
+  mockHuggingFaceMirror.value = ''
 })
 
 describe('fetchModelMetadata', () => {
@@ -715,3 +733,122 @@ describe('downloadModel', () => {
     })
   })
 })
+
+describe('resolveHuggingFaceUrl', () => {
+  it('returns the URL unchanged when no mirror is configured', () => {
+    mockHuggingFaceMirror.value = ''
+    expect(
+      resolveHuggingFaceUrl('https://huggingface.co/org/model/resolve/main/x')
+    ).toBe('https://huggingface.co/org/model/resolve/main/x')
+  })
+
+  it('returns the URL unchanged when the mirror is whitespace', () => {
+    mockHuggingFaceMirror.value = '   '
+    expect(
+      resolveHuggingFaceUrl('https://huggingface.co/org/model/resolve/main/x')
+    ).toBe('https://huggingface.co/org/model/resolve/main/x')
+  })
+
+  it('rewrites huggingface.co to the configured mirror', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    expect(
+      resolveHuggingFaceUrl('https://huggingface.co/org/model/resolve/main/x')
+    ).toBe('https://hf-mirror.com/org/model/resolve/main/x')
+  })
+
+  it('strips a trailing slash from the mirror', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com/'
+    expect(
+      resolveHuggingFaceUrl('https://huggingface.co/org/model/resolve/main/x')
+    ).toBe('https://hf-mirror.com/org/model/resolve/main/x')
+  })
+
+  it('trims surrounding whitespace from the mirror', () => {
+    mockHuggingFaceMirror.value = '  https://hf-mirror.com  '
+    expect(
+      resolveHuggingFaceUrl('https://huggingface.co/org/model/resolve/main/x')
+    ).toBe('https://hf-mirror.com/org/model/resolve/main/x')
+  })
+
+  it('leaves non-HuggingFace URLs untouched', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    expect(
+      resolveHuggingFaceUrl('https://civitai.com/api/download/models/12345')
+    ).toBe('https://civitai.com/api/download/models/12345')
+  })
+
+  it('does not rewrite URLs whose path merely contains huggingface.co', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    expect(
+      resolveHuggingFaceUrl(
+        'https://example.com/huggingface.co/org/model/resolve/main/x'
+      )
+    ).toBe('https://example.com/huggingface.co/org/model/resolve/main/x')
+  })
+
+  it('routes the Desktop2 download through the mirror', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    const desktopDownloadModel = vi
+      .fn<
+        (url: string, filename: string, directory: string) => Promise<boolean>
+      >()
+      .mockResolvedValue(true)
+    window.__comfyDesktop2 = {
+      isRemote: () => false,
+      downloadModel: desktopDownloadModel
+    }
+
+    downloadModel(
+      {
+        name: 'model.safetensors',
+        url: 'https://huggingface.co/org/model/resolve/main/model.safetensors',
+        directory: 'checkpoints'
+      },
+      {}
+    )
+
+    expect(desktopDownloadModel).toHaveBeenCalledWith(
+      'https://hf-mirror.com/org/model/resolve/main/model.safetensors',
+      'model.safetensors',
+      'checkpoints'
+    )
+  })
+
+  it('routes the Electron download store through the mirror', () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    mockIsDesktop.value = true
+
+    downloadModel(
+      {
+        name: 'model.safetensors',
+        url: 'https://huggingface.co/org/model/resolve/main/model.safetensors',
+        directory: 'checkpoints'
+      },
+      { checkpoints: ['/models/checkpoints'] }
+    )
+
+    expect(mockStartDownload).toHaveBeenCalledWith({
+      url: 'https://hf-mirror.com/org/model/resolve/main/model.safetensors',
+      savePath: '/models/checkpoints',
+      filename: 'model.safetensors'
+    })
+  })
+
+  it('routes the file-size HEAD probe through the mirror', async () => {
+    mockHuggingFaceMirror.value = 'https://hf-mirror.com'
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      headers: new Headers({ 'content-length': '42' })
+    })
+
+    const metadata = await fetchModelMetadata(
+      'https://huggingface.co/org/model/resolve/main/probe.safetensors'
+    )
+
+    expect(metadata.fileSize).toBe(42)
+    expect(fetchMock).toHaveBeenCalledWith(
+      'https://hf-mirror.com/org/model/resolve/main/probe.safetensors',
+      { method: 'HEAD' }
+    )
+  })
+})

src/platform/missingModel/missingModelDownload.ts

@@ -1,5 +1,6 @@
 import { downloadUrlToHfRepoUrl, isCivitaiModelUrl } from '@/utils/formatUtil'
 import { isDesktop } from '@/platform/distribution/types'
+import { useSettingStore } from '@/platform/settings/settingStore'
 import { useElectronDownloadStore } from '@/stores/electronDownloadStore'
 import { useSidebarTabStore } from '@/stores/workspace/sidebarTabStore'
 import type { ComfyDesktop2Bridge } from '@/types'
@@ -37,6 +38,25 @@ function isModelUrlAllowlisted(url: string): boolean {
 
 const MODEL_LIBRARY_TAB_ID = 'model-library'
 
+export const HUGGINGFACE_MIRROR_SETTING_ID =
+  'Comfy.ModelLibrary.HuggingFaceMirror'
+
+/**
+ * Rewrites `huggingface.co` URLs to a user-configured mirror so that
+ * users behind networks blocking `huggingface.co` can still use the
+ * "Download All missing models" button. Empty/missing setting returns
+ * the URL unchanged. Non-HuggingFace URLs are returned unchanged.
+ */
+export function resolveHuggingFaceUrl(url: string): string {
+  const mirror = useSettingStore()
+    .get(HUGGINGFACE_MIRROR_SETTING_ID)
+    ?.trim()
+    .replace(/\/+$/, '')
+  if (!mirror) return url
+  if (!hasHuggingFaceHost(url)) return url
+  return url.replace('https://huggingface.co', mirror)
+}
+
 export interface ModelWithUrl {
   name: string
   url: string
@@ -48,7 +68,8 @@ async function startDesktop2ModelDownload(
   model: ModelWithUrl
 ): Promise<void> {
   try {
-    await bridge.downloadModel?.(model.url, model.name, model.directory)
+    const url = resolveHuggingFaceUrl(model.url)
+    await bridge.downloadModel?.(url, model.name, model.directory)
   } catch (error: unknown) {
     console.error('Failed to start Desktop2 model download:', error)
   }
@@ -139,7 +160,7 @@ export function downloadModel(
   if (modelPaths?.[0]) {
     useSidebarTabStore().activeSidebarTabId = MODEL_LIBRARY_TAB_ID
     void useElectronDownloadStore().start({
-      url: model.url,
+      url: resolveHuggingFaceUrl(model.url),
       savePath: modelPaths[0],
       filename: model.name
     })
@@ -225,7 +246,7 @@ const HUGGING_FACE_GATED_ERROR_CODE = 'GatedRepo'
 async function fetchHeadMetadata(url: string): Promise<MetadataFetchResult> {
   try {
     // Deliberately uncredentialed HEADs prevent re-checks from clearing gating.
-    const response = await fetch(url, { method: 'HEAD' })
+    const response = await fetch(resolveHuggingFaceUrl(url), { method: 'HEAD' })
     if (!response.ok) {
       if (
         isTrustedHuggingFaceUrl(url) &&

src/platform/settings/constants/coreSettings.ts

@@ -437,6 +437,14 @@ export const CORE_SETTINGS: SettingParams[] = [
     options: ['filename', 'title'],
     defaultValue: 'title'
   },
+  {
+    id: 'Comfy.ModelLibrary.HuggingFaceMirror',
+    name: 'HuggingFace mirror URL',
+    tooltip:
+      'Optional mirror for Hugging Face model downloads. Leave empty to use the official huggingface.co (e.g. https://hf-mirror.com). Applies to the "Download All missing models" workflow button and the file-size probe shown in the UI.',
+    type: 'text',
+    defaultValue: ''
+  },
   {
     id: 'Comfy.Locale',
     name: 'Language',

src/schemas/apiSchema.ts

@@ -343,6 +343,7 @@ const zSettings = z.object({
   'Comfy.LinkRelease.Action': zLinkReleaseTriggerAction,
   'Comfy.LinkRelease.ActionShift': zLinkReleaseTriggerAction,
   'Comfy.ModelLibrary.AutoLoadAll': z.boolean(),
+  'Comfy.ModelLibrary.HuggingFaceMirror': z.string(),
   'Comfy.ModelLibrary.NameFormat': z.enum(['filename', 'title']),
   'Comfy.NodeSearchBoxImpl.NodePreview': z.boolean(),
   'Comfy.NodeSearchBoxImpl.FollowCursor': z.boolean(),