diff --git a/__init__.py b/__init__.py index d4adf6d..acfb0f8 100644 --- a/__init__.py +++ b/__init__.py @@ -780,6 +780,28 @@ class ModelResolverExtension: "diffusion_models": "diffusion_models", "unet": "diffusion_models", "upscale_models": "upscale_models", + "upscale_model": "upscale_models", + "latent_upscale_model": "latent_upscale_models", + "latent_upscale_models": "latent_upscale_models", + "style_model": "style_models", + "style_models": "style_models", + "gligen": "gligen", + "diffusers": "diffusers", + "vae_approx": "vae_approx", + "audio_encoder": "audio_encoders", + "audio_encoders": "audio_encoders", + "background_removal": "background_removal", + "background_removal_model": "background_removal", + "frame_interpolation": "frame_interpolation", + "frame_interpolation_model": "frame_interpolation", + "geometry_estimation": "geometry_estimation", + "geometry_estimation_model": "geometry_estimation", + "detection": "detection", + "model_patch": "model_patches", + "model_patches": "model_patches", + "photomaker": "photomaker", + "optical_flow": "optical_flow", + "optical_flow_model": "optical_flow", } folder_type = category_map.get( category.lower(), category.lower() @@ -1906,7 +1928,9 @@ class ModelResolverExtension: async def get_directories(request): """Get available model directories.""" try: - categories = [ + import folder_paths + + preferred_categories = [ "checkpoints", "loras", "vae", @@ -1920,6 +1944,20 @@ class ModelResolverExtension: "ipadapter", "sams", ] + skip_categories = {"custom_nodes", "configs"} + categories = [] + for cat in [ + *preferred_categories, + *folder_paths.folder_names_and_paths.keys(), + ]: + normalized_cat = normalize_download_category(cat) + if ( + not normalized_cat + or normalized_cat in skip_categories + or normalized_cat in categories + ): + continue + categories.append(normalized_cat) directories = {} for cat in categories: diff --git a/core/downloader.py b/core/downloader.py index 6bc5879..915caec 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -58,6 +58,25 @@ CATEGORY_MAP = { "latent_upscale_models": "latent_upscale_models", "text_encoders": "text_encoders", "text_encoder": "text_encoders", + "style_model": "style_models", + "style_models": "style_models", + "gligen": "gligen", + "diffusers": "diffusers", + "vae_approx": "vae_approx", + "audio_encoder": "audio_encoders", + "audio_encoders": "audio_encoders", + "background_removal": "background_removal", + "background_removal_model": "background_removal", + "frame_interpolation": "frame_interpolation", + "frame_interpolation_model": "frame_interpolation", + "geometry_estimation": "geometry_estimation", + "geometry_estimation_model": "geometry_estimation", + "detection": "detection", + "model_patch": "model_patches", + "model_patches": "model_patches", + "photomaker": "photomaker", + "optical_flow": "optical_flow", + "optical_flow_model": "optical_flow", "ipadapter": "ipadapter", "ip_adapter": "ipadapter", # Legacy/unsupported categories map to checkpoints as fallback diff --git a/core/scanner.py b/core/scanner.py index 8ec5493..db382e3 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -103,6 +103,23 @@ def scan_directory( # Skip hidden directories dirs[:] = [d for d in dirs if not d.startswith(".")] + if "folder" in (extensions or set()) and root != base_directory: + if category != "diffusers" or "model_index.json" in files: + try: + relative_path = os.path.relpath(root, base_directory) + except ValueError: + relative_path = os.path.basename(root) + + models.append( + { + "filename": os.path.basename(root), + "path": root, + "relative_path": relative_path, + "category": category, + "base_directory": base_directory, + } + ) + for filename in files: # Check if file has a model extension file_ext = os.path.splitext(filename)[1].lower() diff --git a/core/settings.py b/core/settings.py index cab9271..ab538aa 100644 --- a/core/settings.py +++ b/core/settings.py @@ -32,6 +32,27 @@ CATEGORY_MAP = { "upscaler": "upscale_models", "upscale_model": "upscale_models", "upscale_models": "upscale_models", + "latent_upscale_model": "latent_upscale_models", + "latent_upscale_models": "latent_upscale_models", + "style_model": "style_models", + "style_models": "style_models", + "gligen": "gligen", + "diffusers": "diffusers", + "vae_approx": "vae_approx", + "audio_encoder": "audio_encoders", + "audio_encoders": "audio_encoders", + "background_removal": "background_removal", + "background_removal_model": "background_removal", + "frame_interpolation": "frame_interpolation", + "frame_interpolation_model": "frame_interpolation", + "geometry_estimation": "geometry_estimation", + "geometry_estimation_model": "geometry_estimation", + "detection": "detection", + "model_patch": "model_patches", + "model_patches": "model_patches", + "photomaker": "photomaker", + "optical_flow": "optical_flow", + "optical_flow_model": "optical_flow", } DEFAULT_DOWNLOAD_PATH_TEMPLATES: Dict[str, str] = { diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 6959918..0648699 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -45,11 +45,13 @@ URN_TYPE_MAP = { "vae": "vae", "upscaler": "upscale_models", "upscale_model": "upscale_models", + "latent_upscale_model": "latent_upscale_models", "embedding": "embeddings", "hypernetwork": "hypernetworks", "controlnet": "controlnet", "clip": "text_encoders", "clip_vision": "clip_vision", + "diffusers": "diffusers", } URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") @@ -60,16 +62,23 @@ URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") NODE_TYPE_TO_CATEGORY_HINTS = { "CheckpointLoaderSimple": "checkpoints", "CheckpointLoader": "checkpoints", + "DiffusersLoader": "diffusers", "unCLIPCheckpointLoader": "checkpoints", "ImageOnlyCheckpointLoader": "checkpoints", "VAELoader": "vae", "VAELoaderKJ": "vae", "LoraLoader": "loras", "LoraLoaderModelOnly": "loras", + "LoraLoaderBypass": "loras", + "LoraLoaderBypassModelOnly": "loras", "LoraLoaderV2": "loras", "Lora Loader (LoraManager)": "loras", # LoraManager custom node "Lora Stacker (LoraManager)": "loras", # LoraManager Stacker node "Power Lora Loader (rgthree)": "loras", # rgthree's Power Lora Loader + "CreateHookLora": "loras", + "CreateHookLoraModelOnly": "loras", + "CreateHookModelAsLora": "checkpoints", + "CreateHookModelAsLoraModelOnly": "checkpoints", "UNETLoader": "diffusion_models", "LoaderGGUF": "diffusion_models", "LoaderGGUFAdvanced": "diffusion_models", @@ -82,16 +91,32 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "ClipLoaderGGUF": "text_encoders", "DualCLIPLoaderGGUF": "text_encoders", "DualClipLoaderGGUF": "text_encoders", + "TripleCLIPLoader": "text_encoders", + "TripleClipLoader": "text_encoders", "TripleCLIPLoaderGGUF": "text_encoders", "TripleClipLoaderGGUF": "text_encoders", + "QuadrupleCLIPLoader": "text_encoders", + "QuadrupleClipLoader": "text_encoders", "QuadrupleCLIPLoaderGGUF": "text_encoders", "QuadrupleClipLoaderGGUF": "text_encoders", "ControlNetLoader": "controlnet", "DiffControlNetLoader": "controlnet", "ControlNetLoaderAdvanced": "controlnet", + "ACN_ControlNetLoaderAdvanced": "controlnet", + "ACN_DiffControlNetLoaderAdvanced": "controlnet", "CLIPVisionLoader": "clip_vision", "StyleModelLoader": "style_models", + "GLIGENLoader": "gligen", "UpscaleModelLoader": "upscale_models", + "AudioEncoderLoader": "audio_encoders", + "LoadBackgroundRemovalModel": "background_removal", + "LoadDA3Model": "geometry_estimation", + "FrameInterpolationModelLoader": "frame_interpolation", + "LoadMediaPipeFaceLandmarker": "detection", + "ModelPatchLoader": "model_patches", + "LoadMoGeModel": "geometry_estimation", + "PhotoMakerLoader": "photomaker", + "OpticalFlowLoader": "optical_flow", "HypernetworkLoader": "hypernetworks", "EmbeddingLoader": "embeddings", # LTX-Video nodes @@ -106,12 +131,19 @@ NODE_TYPE_TO_CATEGORY_HINTS = { NODE_TYPE_MODEL_WIDGET_CATEGORIES = { "CheckpointLoaderSimple": {0: "checkpoints"}, "CheckpointLoader": {1: "checkpoints"}, + "DiffusersLoader": {0: "diffusers"}, "unCLIPCheckpointLoader": {0: "checkpoints"}, "ImageOnlyCheckpointLoader": {0: "checkpoints"}, "VAELoader": {0: "vae"}, "VAELoaderKJ": {0: "vae"}, "LoraLoader": {0: "loras"}, "LoraLoaderModelOnly": {0: "loras"}, + "LoraLoaderBypass": {0: "loras"}, + "LoraLoaderBypassModelOnly": {0: "loras"}, + "CreateHookLora": {0: "loras"}, + "CreateHookLoraModelOnly": {0: "loras"}, + "CreateHookModelAsLora": {0: "checkpoints"}, + "CreateHookModelAsLoraModelOnly": {0: "checkpoints"}, "UNETLoader": {0: "diffusion_models"}, "LoaderGGUF": {0: "diffusion_models"}, "LoaderGGUFAdvanced": {0: "diffusion_models"}, @@ -124,6 +156,16 @@ NODE_TYPE_MODEL_WIDGET_CATEGORIES = { "ClipLoaderGGUF": {0: "text_encoders"}, "DualCLIPLoaderGGUF": {0: "text_encoders", 1: "text_encoders"}, "DualClipLoaderGGUF": {0: "text_encoders", 1: "text_encoders"}, + "TripleCLIPLoader": { + 0: "text_encoders", + 1: "text_encoders", + 2: "text_encoders", + }, + "TripleClipLoader": { + 0: "text_encoders", + 1: "text_encoders", + 2: "text_encoders", + }, "TripleCLIPLoaderGGUF": { 0: "text_encoders", 1: "text_encoders", @@ -134,6 +176,18 @@ NODE_TYPE_MODEL_WIDGET_CATEGORIES = { 1: "text_encoders", 2: "text_encoders", }, + "QuadrupleCLIPLoader": { + 0: "text_encoders", + 1: "text_encoders", + 2: "text_encoders", + 3: "text_encoders", + }, + "QuadrupleClipLoader": { + 0: "text_encoders", + 1: "text_encoders", + 2: "text_encoders", + 3: "text_encoders", + }, "QuadrupleCLIPLoaderGGUF": { 0: "text_encoders", 1: "text_encoders", @@ -149,9 +203,21 @@ NODE_TYPE_MODEL_WIDGET_CATEGORIES = { "ControlNetLoader": {0: "controlnet"}, "DiffControlNetLoader": {0: "controlnet"}, "ControlNetLoaderAdvanced": {0: "controlnet"}, + "ACN_ControlNetLoaderAdvanced": {0: "controlnet"}, + "ACN_DiffControlNetLoaderAdvanced": {0: "controlnet"}, "CLIPVisionLoader": {0: "clip_vision"}, "StyleModelLoader": {0: "style_models"}, + "GLIGENLoader": {0: "gligen"}, "UpscaleModelLoader": {0: "upscale_models"}, + "AudioEncoderLoader": {0: "audio_encoders"}, + "LoadBackgroundRemovalModel": {0: "background_removal"}, + "LoadDA3Model": {0: "geometry_estimation"}, + "FrameInterpolationModelLoader": {0: "frame_interpolation"}, + "LoadMediaPipeFaceLandmarker": {0: "detection"}, + "ModelPatchLoader": {0: "model_patches"}, + "LoadMoGeModel": {0: "geometry_estimation"}, + "PhotoMakerLoader": {0: "photomaker"}, + "OpticalFlowLoader": {0: "optical_flow"}, "HypernetworkLoader": {0: "hypernetworks"}, "EmbeddingLoader": {0: "embeddings"}, "LTXVAudioVAELoader": {0: "checkpoints"}, @@ -176,13 +242,37 @@ MODEL_WIDGET_NAME_TO_CATEGORY = { "clip_name4": "text_encoders", "clip_vision_name": "clip_vision", "lora_name": "loras", + "existing_lora": "loras", "control_net_name": "controlnet", "cnet": "controlnet", "style_model_name": "style_models", "upscale_model_name": "upscale_models", + "gligen_name": "gligen", + "audio_encoder_name": "audio_encoders", + "bg_removal_name": "background_removal", + "photomaker_model_name": "photomaker", + "text_encoder": "text_encoders", "hypernetwork_name": "hypernetworks", } +MODEL_OUTPUT_TYPE_TO_CATEGORY = { + "UPSCALE_MODEL": "upscale_models", + "LATENT_UPSCALE_MODEL": "latent_upscale_models", + "CONTROL_NET": "controlnet", + "CLIP_VISION": "clip_vision", + "STYLE_MODEL": "style_models", + "GLIGEN": "gligen", + "AUDIO_ENCODER": "audio_encoders", + "BACKGROUND_REMOVAL": "background_removal", + "DA3_MODEL": "geometry_estimation", + "MOGE_MODEL": "geometry_estimation", + "INTERP_MODEL": "frame_interpolation", + "FACE_DETECTION_MODEL": "detection", + "MODEL_PATCH": "model_patches", + "PHOTOMAKER": "photomaker", + "OPTICAL_FLOW": "optical_flow", +} + # Keys within dict-type widget values that contain model file references. # Some nodes (e.g. rgthree Power Lora Loader) store model info as objects like # {"on": true, "lora": "name.safetensors", "strength": 1.0} inside widgets_values. @@ -202,8 +292,18 @@ NESTED_MODEL_KEYS = { "model_name": "diffusion_models", "unet_name": "diffusion_models", "gguf_name": "diffusion_models", + "gligen_name": "gligen", + "audio_encoder_name": "audio_encoders", + "bg_removal_name": "background_removal", + "photomaker_model_name": "photomaker", + "text_encoder": "text_encoders", } + +def normalize_widget_name(value: Any) -> str: + return re.sub(r"[_\s-]+", "_", str(value or "").strip().lower()).strip("_") + + MODEL_WIDGET_PLACEHOLDERS = { "", "none", @@ -253,7 +353,7 @@ def get_widget_category_hint(node: Dict[str, Any], widget_index: int) -> Optiona candidates.append(str(widget).strip()) for candidate in candidates: - category = MODEL_WIDGET_NAME_TO_CATEGORY.get(candidate) + category = MODEL_WIDGET_NAME_TO_CATEGORY.get(normalize_widget_name(candidate)) if category: return category @@ -267,6 +367,23 @@ def get_node_model_widget_category_hint( return NODE_TYPE_MODEL_WIDGET_CATEGORIES.get(node_type, {}).get(widget_index) +def get_node_output_category_hint(node: Dict[str, Any]) -> Optional[str]: + """Return a model category hint from strongly typed loader outputs.""" + outputs = node.get("outputs", []) + if not isinstance(outputs, list): + return None + + for output in outputs: + if not isinstance(output, dict): + continue + for key in ("type", "name", "label"): + token = str(output.get(key, "") or "").strip().upper() + category = MODEL_OUTPUT_TYPE_TO_CATEGORY.get(token) + if category: + return category + return None + + def is_placeholder_model_value(value: Any) -> bool: if not isinstance(value, str): return False @@ -608,12 +725,21 @@ def get_node_model_info( # Get category hints for this node type category_hint = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type) + output_category_hint = get_node_output_category_hint(node) + has_single_widget_value = len(widgets_values) == 1 # For each widget value, check if it looks like a model file or URN for idx, value in enumerate(widgets_values): widget_category_hint = get_widget_category_hint(node, idx) indexed_category_hint = get_node_model_widget_category_hint(node_type, idx) - model_widget_category_hint = widget_category_hint or indexed_category_hint + output_widget_category_hint = ( + output_category_hint + if (indexed_category_hint or widget_category_hint or has_single_widget_value) + else None + ) + model_widget_category_hint = ( + indexed_category_hint or output_widget_category_hint or widget_category_hint + ) effective_category_hint = model_widget_category_hint or category_hint categories_to_try_for_widget = ( [effective_category_hint] if effective_category_hint else None @@ -689,8 +815,7 @@ def get_node_model_info( if urn_match: base, typ, provider, model_id, version_id = urn_match.groups() category = ( - category_hint - or widget_category_hint + effective_category_hint or URN_TYPE_MAP.get(typ.lower(), "unknown") ) diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index 3d69886..39cec6d 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -117,6 +117,15 @@ export const downloadTargetMethods = { text_encoder: 'text_encoders', text_encoders: 'text_encoders', ip_adapter: 'ipadapter', + upscale_model: 'upscale_models', + latent_upscale_model: 'latent_upscale_models', + style_model: 'style_models', + audio_encoder: 'audio_encoders', + model_patch: 'model_patches', + background_removal_model: 'background_removal', + frame_interpolation_model: 'frame_interpolation', + geometry_estimation_model: 'geometry_estimation', + optical_flow_model: 'optical_flow', 'default': 'upscale_models' }; return categoryMap[token] || token || 'checkpoints'; @@ -136,7 +145,19 @@ export const downloadTargetMethods = { 'text_encoders': 'text encoders', 'clip': 'clip', 'clip_vision': 'clip_vision', - 'hypernetworks': 'hypernetwork' + 'hypernetworks': 'hypernetwork', + 'style_models': 'style model', + 'gligen': 'GLIGEN', + 'diffusers': 'Diffusers', + 'vae_approx': 'TAESD / VAE approx', + 'audio_encoders': 'audio encoder', + 'background_removal': 'background removal', + 'frame_interpolation': 'frame interpolation', + 'geometry_estimation': 'geometry estimation', + 'detection': 'detection', + 'model_patches': 'model patch', + 'photomaker': 'PhotoMaker', + 'optical_flow': 'optical flow' }; return displayNames[category] || category || 'unknown'; }, @@ -196,6 +217,18 @@ export const downloadTargetMethods = { latent_upscale_models: 'latent', latent_upscale_model: 'latent', latent: 'latent', + gligen: 'conditioning', + diffusers: 'model', + audio_encoders: 'clip', + audio_encoder: 'clip', + background_removal: 'image', + frame_interpolation: 'image', + geometry_estimation: 'image', + detection: 'controlnet', + model_patches: 'model', + model_patch: 'model', + photomaker: 'clip-vision', + optical_flow: 'image', image: 'image', images: 'image', mask: 'mask', @@ -224,7 +257,19 @@ export const downloadTargetMethods = { 'text_encoders': 'text_encoders', 'clip': 'clip', 'clip_vision': 'clip_vision', - 'hypernetworks': 'hypernetwork' + 'hypernetworks': 'hypernetwork', + 'style_models': 'style_model', + 'gligen': 'gligen', + 'diffusers': 'diffusers', + 'vae_approx': 'vae_approx', + 'audio_encoders': 'audio_encoder', + 'background_removal': 'background_removal', + 'frame_interpolation': 'frame_interpolation', + 'geometry_estimation': 'geometry_estimation', + 'detection': 'detection', + 'model_patches': 'model_patch', + 'photomaker': 'photomaker', + 'optical_flow': 'optical_flow' }; return tokenNames[category] || category || 'unknown'; }, @@ -240,6 +285,18 @@ export const downloadTargetMethods = { 'upscale_models', 'controlnet', 'clip_vision', + 'style_models', + 'gligen', + 'diffusers', + 'latent_upscale_models', + 'audio_encoders', + 'background_removal', + 'frame_interpolation', + 'geometry_estimation', + 'detection', + 'model_patches', + 'photomaker', + 'optical_flow', 'ipadapter', 'sams' ]; diff --git a/web/resolver/shell/workflow_state_methods.js b/web/resolver/shell/workflow_state_methods.js index f3e8748..5a0750d 100644 --- a/web/resolver/shell/workflow_state_methods.js +++ b/web/resolver/shell/workflow_state_methods.js @@ -33,10 +33,17 @@ export const workflowStateMethods = { 'control_net_name', 'cnet', 'model_name', + 'model_path', 'unet_name', 'gguf_name', 'upscale_model_name', 'style_model_name', + 'gligen_name', + 'audio_encoder_name', + 'bg_removal_name', + 'photomaker_model_name', + 'text_encoder', + 'existing_lora', 'hypernetwork_name', 'clip_name3', 'clip_name4' @@ -62,12 +69,19 @@ export const workflowStateMethods = { const modelWidgetIndicesByNodeType = { CheckpointLoaderSimple: [0], CheckpointLoader: [1], + DiffusersLoader: [0], unCLIPCheckpointLoader: [0], ImageOnlyCheckpointLoader: [0], VAELoader: [0], VAELoaderKJ: [0], LoraLoader: [0], LoraLoaderModelOnly: [0], + LoraLoaderBypass: [0], + LoraLoaderBypassModelOnly: [0], + CreateHookLora: [0], + CreateHookLoraModelOnly: [0], + CreateHookModelAsLora: [0], + CreateHookModelAsLoraModelOnly: [0], UNETLoader: [0], LoaderGGUF: [0], LoaderGGUFAdvanced: [0], @@ -80,16 +94,32 @@ export const workflowStateMethods = { ClipLoaderGGUF: [0], DualCLIPLoaderGGUF: [0, 1], DualClipLoaderGGUF: [0, 1], + TripleCLIPLoader: [0, 1, 2], + TripleClipLoader: [0, 1, 2], TripleCLIPLoaderGGUF: [0, 1, 2], TripleClipLoaderGGUF: [0, 1, 2], + QuadrupleCLIPLoader: [0, 1, 2, 3], + QuadrupleClipLoader: [0, 1, 2, 3], QuadrupleCLIPLoaderGGUF: [0, 1, 2, 3], QuadrupleClipLoaderGGUF: [0, 1, 2, 3], ControlNetLoader: [0], DiffControlNetLoader: [0], ControlNetLoaderAdvanced: [0], + ACN_ControlNetLoaderAdvanced: [0], + ACN_DiffControlNetLoaderAdvanced: [0], CLIPVisionLoader: [0], StyleModelLoader: [0], + GLIGENLoader: [0], UpscaleModelLoader: [0], + AudioEncoderLoader: [0], + LoadBackgroundRemovalModel: [0], + LoadDA3Model: [0], + FrameInterpolationModelLoader: [0], + LoadMediaPipeFaceLandmarker: [0], + ModelPatchLoader: [0], + LoadMoGeModel: [0], + PhotoMakerLoader: [0], + OpticalFlowLoader: [0], HypernetworkLoader: [0], EmbeddingLoader: [0], LTXVAudioVAELoader: [0],