diff --git a/src/stores/modelToNodeStore.test.ts b/src/stores/modelToNodeStore.test.ts index de57a5aef..ceaeab6c4 100644 --- a/src/stores/modelToNodeStore.test.ts +++ b/src/stores/modelToNodeStore.test.ts @@ -27,7 +27,19 @@ const EXPECTED_DEFAULT_TYPES = [ 'chatterbox/chatterbox', 'chatterbox/chatterbox_turbo', 'chatterbox/chatterbox_multilingual', - 'chatterbox/chatterbox_vc' + 'chatterbox/chatterbox_vc', + 'latent_upscale_models', + 'sam2', + 'sams', + 'ultralytics', + 'depthanything', + 'ipadapter', + 'segformer_b2_clothes', + 'segformer_b3_clothes', + 'segformer_b3_fashion', + 'nlf', + 'FlashVSR', + 'FlashVSR-v1.1' ] as const type NodeDefStoreType = ReturnType @@ -69,7 +81,17 @@ const MOCK_NODE_NAMES = [ 'FL_ChatterboxTTS', 'FL_ChatterboxTurboTTS', 'FL_ChatterboxMultilingualTTS', - 'FL_ChatterboxVC' + 'FL_ChatterboxVC', + // New extension node mappings + 'LatentUpscaleModelLoader', + 'DownloadAndLoadSAM2Model', + 'SAMLoader', + 'UltralyticsDetectorProvider', + 'DownloadAndLoadDepthAnythingV2Model', + 'IPAdapterModelLoader', + 'LS_LoadSegformerModel', + 'LoadNLFModel', + 'FlashVSRNode' ] as const const mockNodeDefsByName = Object.fromEntries( @@ -173,6 +195,79 @@ describe('useModelToNodeStore', () => { expect(provider?.nodeDef?.name).toBe('FL_ChatterboxVC') expect(provider?.key).toBe('') }) + + it('should return provider for new extension model types', () => { + const modelToNodeStore = useModelToNodeStore() + modelToNodeStore.registerDefaults() + + // SAM2 + const sam2Provider = modelToNodeStore.getNodeProvider('sam2') + expect(sam2Provider?.nodeDef?.name).toBe('DownloadAndLoadSAM2Model') + expect(sam2Provider?.key).toBe('model') + + // SAMLoader (original SAM) + const samsProvider = modelToNodeStore.getNodeProvider('sams') + expect(samsProvider?.nodeDef?.name).toBe('SAMLoader') + expect(samsProvider?.key).toBe('model_name') + + // IP-Adapter + const ipadapterProvider = modelToNodeStore.getNodeProvider('ipadapter') + expect(ipadapterProvider?.nodeDef?.name).toBe('IPAdapterModelLoader') + expect(ipadapterProvider?.key).toBe('ipadapter_file') + + // DepthAnything + const depthProvider = modelToNodeStore.getNodeProvider('depthanything') + expect(depthProvider?.nodeDef?.name).toBe( + 'DownloadAndLoadDepthAnythingV2Model' + ) + expect(depthProvider?.key).toBe('model') + }) + + it('should use hierarchical fallback for ultralytics subcategories', () => { + const modelToNodeStore = useModelToNodeStore() + modelToNodeStore.registerDefaults() + + // ultralytics/bbox should fall back to ultralytics + const bboxProvider = modelToNodeStore.getNodeProvider('ultralytics/bbox') + expect(bboxProvider?.nodeDef?.name).toBe('UltralyticsDetectorProvider') + expect(bboxProvider?.key).toBe('model_name') + + // ultralytics/segm should also fall back to ultralytics + const segmProvider = modelToNodeStore.getNodeProvider('ultralytics/segm') + expect(segmProvider?.nodeDef?.name).toBe('UltralyticsDetectorProvider') + }) + + it('should return provider for FlashVSR nodes with empty key (auto-load)', () => { + const modelToNodeStore = useModelToNodeStore() + modelToNodeStore.registerDefaults() + + const flashVSRProvider = modelToNodeStore.getNodeProvider('FlashVSR') + expect(flashVSRProvider?.nodeDef?.name).toBe('FlashVSRNode') + expect(flashVSRProvider?.key).toBe('') + + const flashVSR11Provider = + modelToNodeStore.getNodeProvider('FlashVSR-v1.1') + expect(flashVSR11Provider?.nodeDef?.name).toBe('FlashVSRNode') + expect(flashVSR11Provider?.key).toBe('') + }) + + it('should return provider for segformer models', () => { + const modelToNodeStore = useModelToNodeStore() + modelToNodeStore.registerDefaults() + + const segformerB2Provider = modelToNodeStore.getNodeProvider( + 'segformer_b2_clothes' + ) + expect(segformerB2Provider?.nodeDef?.name).toBe('LS_LoadSegformerModel') + expect(segformerB2Provider?.key).toBe('model_name') + + const segformerB3FashionProvider = modelToNodeStore.getNodeProvider( + 'segformer_b3_fashion' + ) + expect(segformerB3FashionProvider?.nodeDef?.name).toBe( + 'LS_LoadSegformerModel' + ) + }) }) describe('getAllNodeProviders', () => { diff --git a/src/stores/modelToNodeStore.ts b/src/stores/modelToNodeStore.ts index a17aee9a8..67e7080ad 100644 --- a/src/stores/modelToNodeStore.ts +++ b/src/stores/modelToNodeStore.ts @@ -192,6 +192,44 @@ export const useModelToNodeStore = defineStore('modelToNode', () => { '' ) quickRegister('chatterbox/chatterbox_vc', 'FL_ChatterboxVC', '') + + // Latent upscale models (ComfyUI core - nodes_hunyuan.py) + quickRegister( + 'latent_upscale_models', + 'LatentUpscaleModelLoader', + 'model_name' + ) + + // SAM/SAM2 segmentation models (comfyui-segment-anything-2, comfyui-impact-pack) + quickRegister('sam2', 'DownloadAndLoadSAM2Model', 'model') + quickRegister('sams', 'SAMLoader', 'model_name') + + // Ultralytics detection models (comfyui-impact-subpack) + // Note: ultralytics/bbox and ultralytics/segm fall back to this via hierarchical lookup + quickRegister('ultralytics', 'UltralyticsDetectorProvider', 'model_name') + + // DepthAnything models (comfyui-depthanythingv2) + quickRegister( + 'depthanything', + 'DownloadAndLoadDepthAnythingV2Model', + 'model' + ) + + // IP-Adapter models (comfyui_ipadapter_plus) + quickRegister('ipadapter', 'IPAdapterModelLoader', 'ipadapter_file') + + // Segformer clothing/fashion segmentation models (comfyui_layerstyle) + quickRegister('segformer_b2_clothes', 'LS_LoadSegformerModel', 'model_name') + quickRegister('segformer_b3_clothes', 'LS_LoadSegformerModel', 'model_name') + quickRegister('segformer_b3_fashion', 'LS_LoadSegformerModel', 'model_name') + + // NLF pose estimation models (ComfyUI-WanVideoWrapper) + quickRegister('nlf', 'LoadNLFModel', 'nlf_model') + + // FlashVSR video super-resolution (ComfyUI-FlashVSR_Ultra_Fast) + // Empty key means the node auto-loads models without a widget selector + quickRegister('FlashVSR', 'FlashVSRNode', '') + quickRegister('FlashVSR-v1.1', 'FlashVSRNode', '') } return {