diff --git a/core/routes/loaded_models.py b/core/routes/loaded_models.py index 2523fec..21ebaa1 100644 --- a/core/routes/loaded_models.py +++ b/core/routes/loaded_models.py @@ -1,479 +1,21 @@ """Loaded model route registration.""" +from ..services.loaded_models_service import LoadedModelsService from .context import RouteContext def register_loaded_model_routes(context: RouteContext): - adapt_custom_node_loaded_model = context.get('adapt_custom_node_loaded_model') - asyncio = context.get('asyncio') - get_filename_from_path = context.get('get_filename_from_path') - get_progress_response = context.get('get_progress_response') - get_workflow_model_inventory = context.get('get_workflow_model_inventory') - json_api_endpoint = context.get('json_api_endpoint') - routes = context.get('routes') - self = context.get('self') - web = context.get('web') + """Register loaded model routes and delegate behavior to a service.""" + json_api_endpoint = context.get("json_api_endpoint") + routes = context.get("routes") + service = LoadedModelsService(context) @routes.post("/model_resolver/loaded") @json_api_endpoint("get_loaded_models") async def get_loaded_models(request): - """Get all currently loaded models in the workflow.""" - data = await request.json() - workflow_json = data.get("workflow") - loaded_id = str( - data.get("loaded_id") or data.get("progress_id") or "" - ).strip() - - if not workflow_json: - return web.json_response( - {"error": "Workflow JSON is required"}, status=400 - ) - if not isinstance(workflow_json, dict): - return web.json_response( - {"error": "Workflow JSON must be an object"}, status=400 - ) - - def update_loaded_progress(*args, **kwargs): - self._update_loaded_progress(loaded_id, *args, **kwargs) - - def get_workflow_node_count(): - node_count = 0 - nodes = workflow_json.get("nodes", []) - if isinstance(nodes, list): - node_count += len(nodes) - - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] - ) - if isinstance(subgraphs, list): - for subgraph in subgraphs: - if not isinstance(subgraph, dict): - continue - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - node_count += len(subgraph_nodes) - return node_count - - def interpolate_percent(start, end, current, total): - if not total: - return start - try: - ratio = max(0.0, min(1.0, float(current) / float(total))) - except (TypeError, ValueError, ZeroDivisionError): - ratio = 0.0 - return start + ((end - start) * ratio) - - def build_loaded_models_response(): - from ..workflow_analyzer import ( - URN_TYPE_MAP, - get_lora_model_strength, - ) - - update_loaded_progress( - "scanning", - "Scanning local model index...", - percent=5, - ) - - workflow_node_count = get_workflow_node_count() - - def update_workflow_analysis_progress(payload): - self._update_workflow_analysis_progress( - loaded_id, - workflow_node_count, - interpolate_percent, - payload, - start_percent=5, - end_percent=45, - ) - - inventory = get_workflow_model_inventory( - workflow_json, - progress_callback=update_workflow_analysis_progress, - ) - available_models = inventory["available_models"] - all_model_refs = inventory["model_refs"] - - # Create lookup for full paths by filename (with and without extension) - path_by_filename = {} - total_local_models = len(available_models) - update_loaded_progress( - "indexing", - f"Indexing {total_local_models} local models...", - percent=45, - current=0, - total=total_local_models, - ) - for index, model_info in enumerate(available_models, start=1): - rel_path = model_info.get("relative_path", "") - if rel_path: - filename = get_filename_from_path(rel_path) - path_by_filename[filename] = model_info.get("path") - # Also add without extension for matching (simple approach) - if "." in filename: - filename_no_ext = filename.rsplit(".", 1)[0] - if filename_no_ext not in path_by_filename: - path_by_filename[filename_no_ext] = model_info.get("path") - # Add the full relative path as key too - path_by_filename[rel_path] = model_info.get("path") - - if ( - total_local_models - and (index == total_local_models or index % 250 == 0) - ): - update_loaded_progress( - "indexing", - f"Indexing local model {index} of {total_local_models}", - percent=interpolate_percent( - 45, 55, index, total_local_models - ), - current=index, - total=total_local_models, - ) - - # Also use folder_paths.get_full_path() to get paths - try: - import folder_paths - except Exception: - folder_paths = None - - folder_categories = [ - "loras", - "checkpoints", - "vae", - "controlnet", - "upscale_models", - ] - if folder_paths: - for index, category_name in enumerate(folder_categories, start=1): - update_loaded_progress( - "indexing", - f"Reading {category_name} model list...", - percent=interpolate_percent( - 55, 65, index - 1, len(folder_categories) - ), - current=index - 1, - total=len(folder_categories), - ) - try: - filenames = folder_paths.get_filename_list(category_name) - for filename in filenames: - full_path = folder_paths.get_full_path( - category_name, filename - ) - if ( - full_path - and full_path not in path_by_filename.values() - ): - path_by_filename[filename] = full_path - filename_no_ext = ( - filename.rsplit(".", 1)[0] - if "." in filename - else filename - ) - if filename_no_ext not in path_by_filename: - path_by_filename[filename_no_ext] = full_path - except Exception: - pass - - # Also extract from node.properties.models - workflow_nodes = workflow_json.get("nodes", []) - nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] - node_scopes = { - id(node): { - "subgraph_id": "", - "subgraph_name": "", - "is_top_level": True, - } - for node in nodes - if isinstance(node, dict) - } - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] - ) - for subgraph in subgraphs: - if isinstance(subgraph, dict): - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - nodes.extend(subgraph_nodes) - for node in subgraph_nodes: - if not isinstance(node, dict): - continue - node_scopes[id(node)] = { - "subgraph_id": subgraph.get("id") or "", - "subgraph_name": ( - subgraph.get("name") - or subgraph.get("id") - or "" - ), - "is_top_level": False, - } - - def get_node_scope(node): - return node_scopes.get( - id(node), - { - "subgraph_id": "", - "subgraph_name": "", - "is_top_level": True, - }, - ) - - def node_matches_ref(node, ref): - if not isinstance(node, dict): - return False - if str(node.get("id")) != str(ref.get("node_id")): - return False - scope = get_node_scope(node) - if scope["is_top_level"] != ( - ref.get("is_top_level") is not False - ): - return False - if not scope["is_top_level"]: - return str(scope["subgraph_id"]) == str( - ref.get("subgraph_id") or "" - ) - return True - - # Collect all loaded models with their values - loaded_models = [] - total_refs = len(all_model_refs) - update_loaded_progress( - "building", - "Building loaded models list...", - percent=65, - current=0, - total=total_refs, - ) - - # Process each model reference from analyze_workflow_models - for index, ref in enumerate(all_model_refs, start=1): - if total_refs and (index == total_refs or index % 50 == 0): - update_loaded_progress( - "building", - f"Building loaded model {index} of {total_refs}", - percent=interpolate_percent(65, 94, index, total_refs), - current=index, - total=total_refs, - ) - - original_path = ref.get("original_path", "") - node_id = ref.get("node_id") - widget_index = ref.get("widget_index") - node_type = ref.get("node_type", "") - category = ref.get("category", "unknown") - - # Determine model name and strength - model_name = get_filename_from_path(original_path) - strength = None - - if ( - ref.get("strength") is None - and category in {"lora", "loras"} - ): - for node in nodes: - if node_matches_ref(node, ref): - strength = get_lora_model_strength( - node, - widget_index, - ) - break - - if ref.get("strength") is not None: - strength = ref.get("strength") - - model_name, strength = adapt_custom_node_loaded_model( - ref, - model_name, - strength, - ) - - # Check if model exists locally - exists = ref.get("exists", False) - - # If URN, resolve to display name - if ref.get("is_urn"): - urn = ref.get("urn", {}) - # Use model name from URN as display name - model_name = ( - f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}" - ) - category = urn.get("type", category) - if category in URN_TYPE_MAP: - category = URN_TYPE_MAP[category] - - loaded_models.append( - { - "name": model_name, - "category": category, - "node_id": node_id, - "widget_index": widget_index, - "node_type": node_type, - "node_title": ref.get("node_title", ""), - "subgraph_id": ref.get("subgraph_id") or "", - "subgraph_name": ref.get("subgraph_name") or "", - "is_top_level": ref.get("is_top_level") is not False, - "locate_node_id": ref.get("locate_node_id"), - "locate_node_type": ref.get("locate_node_type", ""), - "locate_node_title": ref.get("locate_node_title", ""), - "locate_subgraph_id": ( - ref.get("locate_subgraph_id") or "" - ), - "locate_subgraph_name": ( - ref.get("locate_subgraph_name") or "" - ), - "locate_is_top_level": ( - ref.get("locate_is_top_level") is not False - ), - "locate_via_promoted_widget": ref.get( - "locate_via_promoted_widget", False - ), - "exists": exists, - "strength": strength, - "original_path": original_path, - "is_urn": ref.get("is_urn", False), - "custom_node_adapter": ref.get( - "custom_node_adapter" - ), - "active": ref.get("active"), - "connected": ref.get("connected", True), - "resolved_path": ( - path_by_filename.get(model_name) - or path_by_filename.get(original_path) - ), - } - ) - - # Also check node.properties.models for embedded models - total_nodes = len(nodes) - for node_index, node in enumerate(nodes, start=1): - if total_nodes and ( - node_index == total_nodes or node_index % 50 == 0 - ): - update_loaded_progress( - "embedded", - "Checking embedded model metadata...", - percent=interpolate_percent( - 94, 98, node_index, total_nodes - ), - current=node_index, - total=total_nodes, - ) - - node_type = node.get("type", "") - properties = node.get("properties", {}) - if not isinstance(properties, dict): - properties = {} - models_list = properties.get("models", []) - if not isinstance(models_list, list): - models_list = [] - - for model_info in models_list: - if isinstance(model_info, dict): - name = model_info.get("name", "") - directory = model_info.get("directory", "") - node_scope = get_node_scope(node) - - if name: - # Check if this model is already in loaded_models - existing = next( - ( - model - for model in loaded_models - if model.get("original_path") == name - ), - None, - ) - if not existing: - loaded_models.append( - { - "name": get_filename_from_path(name), - "category": directory or "checkpoints", - "node_id": node.get("id"), - "widget_index": None, - "node_type": node_type, - "node_title": node.get("title", ""), - "subgraph_id": node_scope["subgraph_id"], - "subgraph_name": node_scope[ - "subgraph_name" - ], - "is_top_level": node_scope[ - "is_top_level" - ], - "exists": True, # Embedded models are loaded - "strength": None, - "original_path": name, - "is_urn": False, - } - ) - - update_loaded_progress( - "finalizing", - "Loaded models ready", - percent=99, - current=len(loaded_models), - total=len(loaded_models), - ) - return { - "loaded_models": loaded_models, - "total": len(loaded_models), - } - - update_loaded_progress( - "starting", - "Preparing loaded model scan...", - percent=0, - status="starting", - current=0, - total=0, - ) - - try: - result = await asyncio.to_thread(build_loaded_models_response) - except Exception as e: - update_loaded_progress( - "error", - str(e), - percent=100, - status="error", - current=0, - total=0, - ) - self.logger.error( - f"Model Resolver loaded models error: {e}", exc_info=True - ) - return web.json_response({"error": str(e)}, status=500) - - update_loaded_progress( - "completed", - "Loaded models ready", - percent=100, - status="completed", - current=result.get("total", 0), - total=result.get("total", 0), - ) - return web.json_response(result) + return await service.get_loaded_models(request) @routes.get("/model_resolver/loaded-progress/{loaded_id}") @json_api_endpoint("loaded-progress") async def get_loaded_models_progress(request): - """Get loaded models inspection progress.""" - return get_progress_response( - self.loaded_progress, - request, - param_name="loaded_id", - not_found_payload={ - "status": "unknown", - "stage": "unknown", - "message": "No loaded models progress available", - "percent": 0, - "current": 0, - "total": 0, - } - ) + return await service.get_loaded_models_progress(request) diff --git a/core/services/loaded_models_service.py b/core/services/loaded_models_service.py new file mode 100644 index 0000000..9391e6a --- /dev/null +++ b/core/services/loaded_models_service.py @@ -0,0 +1,484 @@ +"""Loaded model inspection used by the HTTP route adapter.""" + +from ..routes.context import RouteContext + + +class LoadedModelsService: + """Build the list of models loaded by a workflow and expose its progress.""" + + def __init__(self, context: RouteContext): + self.adapt_custom_node_loaded_model = context.get( + "adapt_custom_node_loaded_model" + ) + self.asyncio = context.get("asyncio") + self.get_filename_from_path = context.get("get_filename_from_path") + self.get_progress_response = context.get("get_progress_response") + self.get_workflow_model_inventory = context.get( + "get_workflow_model_inventory" + ) + self.extension = context.get("self") + self.web = context.get("web") + + async def get_loaded_models(self, request): + """Get all currently loaded models in the workflow.""" + data = await request.json() + workflow_json = data.get("workflow") + loaded_id = str( + data.get("loaded_id") or data.get("progress_id") or "" + ).strip() + + if not workflow_json: + return self.web.json_response( + {"error": "Workflow JSON is required"}, status=400 + ) + if not isinstance(workflow_json, dict): + return self.web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + def update_loaded_progress(*args, **kwargs): + self.extension._update_loaded_progress(loaded_id, *args, **kwargs) + + def get_workflow_node_count(): + node_count = 0 + nodes = workflow_json.get("nodes", []) + if isinstance(nodes, list): + node_count += len(nodes) + + definitions = workflow_json.get("definitions", {}) + subgraphs = ( + definitions.get("subgraphs", []) + if isinstance(definitions, dict) + else [] + ) + if isinstance(subgraphs, list): + for subgraph in subgraphs: + if not isinstance(subgraph, dict): + continue + subgraph_nodes = subgraph.get("nodes", []) + if isinstance(subgraph_nodes, list): + node_count += len(subgraph_nodes) + return node_count + + def interpolate_percent(start, end, current, total): + if not total: + return start + try: + ratio = max(0.0, min(1.0, float(current) / float(total))) + except (TypeError, ValueError, ZeroDivisionError): + ratio = 0.0 + return start + ((end - start) * ratio) + + def build_loaded_models_response(): + from ..workflow_analyzer import ( + URN_TYPE_MAP, + get_lora_model_strength, + ) + + update_loaded_progress( + "scanning", + "Scanning local model index...", + percent=5, + ) + + workflow_node_count = get_workflow_node_count() + + def update_workflow_analysis_progress(payload): + self.extension._update_workflow_analysis_progress( + loaded_id, + workflow_node_count, + interpolate_percent, + payload, + start_percent=5, + end_percent=45, + ) + + inventory = self.get_workflow_model_inventory( + workflow_json, + progress_callback=update_workflow_analysis_progress, + ) + available_models = inventory["available_models"] + all_model_refs = inventory["model_refs"] + + # Create lookup for full paths by filename (with and without extension) + path_by_filename = {} + total_local_models = len(available_models) + update_loaded_progress( + "indexing", + f"Indexing {total_local_models} local models...", + percent=45, + current=0, + total=total_local_models, + ) + for index, model_info in enumerate(available_models, start=1): + rel_path = model_info.get("relative_path", "") + if rel_path: + filename = self.get_filename_from_path(rel_path) + path_by_filename[filename] = model_info.get("path") + # Also add without extension for matching (simple approach) + if "." in filename: + filename_no_ext = filename.rsplit(".", 1)[0] + if filename_no_ext not in path_by_filename: + path_by_filename[filename_no_ext] = model_info.get( + "path" + ) + # Add the full relative path as key too + path_by_filename[rel_path] = model_info.get("path") + + if ( + total_local_models + and (index == total_local_models or index % 250 == 0) + ): + update_loaded_progress( + "indexing", + f"Indexing local model {index} of {total_local_models}", + percent=interpolate_percent( + 45, 55, index, total_local_models + ), + current=index, + total=total_local_models, + ) + + # Also use folder_paths.get_full_path() to get paths + try: + import folder_paths + except Exception: + folder_paths = None + + folder_categories = [ + "loras", + "checkpoints", + "vae", + "controlnet", + "upscale_models", + ] + if folder_paths: + for index, category_name in enumerate(folder_categories, start=1): + update_loaded_progress( + "indexing", + f"Reading {category_name} model list...", + percent=interpolate_percent( + 55, 65, index - 1, len(folder_categories) + ), + current=index - 1, + total=len(folder_categories), + ) + try: + filenames = folder_paths.get_filename_list(category_name) + for filename in filenames: + full_path = folder_paths.get_full_path( + category_name, filename + ) + if ( + full_path + and full_path not in path_by_filename.values() + ): + path_by_filename[filename] = full_path + filename_no_ext = ( + filename.rsplit(".", 1)[0] + if "." in filename + else filename + ) + if filename_no_ext not in path_by_filename: + path_by_filename[filename_no_ext] = full_path + except Exception: + pass + + # Also extract from node.properties.models + workflow_nodes = workflow_json.get("nodes", []) + nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] + node_scopes = { + id(node): { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + } + for node in nodes + if isinstance(node, dict) + } + definitions = workflow_json.get("definitions", {}) + subgraphs = ( + definitions.get("subgraphs", []) + if isinstance(definitions, dict) + else [] + ) + for subgraph in subgraphs: + if isinstance(subgraph, dict): + subgraph_nodes = subgraph.get("nodes", []) + if isinstance(subgraph_nodes, list): + nodes.extend(subgraph_nodes) + for node in subgraph_nodes: + if not isinstance(node, dict): + continue + node_scopes[id(node)] = { + "subgraph_id": subgraph.get("id") or "", + "subgraph_name": ( + subgraph.get("name") + or subgraph.get("id") + or "" + ), + "is_top_level": False, + } + + def get_node_scope(node): + return node_scopes.get( + id(node), + { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + }, + ) + + def node_matches_ref(node, ref): + if not isinstance(node, dict): + return False + if str(node.get("id")) != str(ref.get("node_id")): + return False + scope = get_node_scope(node) + if scope["is_top_level"] != ( + ref.get("is_top_level") is not False + ): + return False + if not scope["is_top_level"]: + return str(scope["subgraph_id"]) == str( + ref.get("subgraph_id") or "" + ) + return True + + # Collect all loaded models with their values + loaded_models = [] + total_refs = len(all_model_refs) + update_loaded_progress( + "building", + "Building loaded models list...", + percent=65, + current=0, + total=total_refs, + ) + + # Process each model reference from analyze_workflow_models + for index, ref in enumerate(all_model_refs, start=1): + if total_refs and ( + index == total_refs or index % 50 == 0 + ): + update_loaded_progress( + "building", + f"Building loaded model {index} of {total_refs}", + percent=interpolate_percent(65, 94, index, total_refs), + current=index, + total=total_refs, + ) + + original_path = ref.get("original_path", "") + node_id = ref.get("node_id") + widget_index = ref.get("widget_index") + node_type = ref.get("node_type", "") + category = ref.get("category", "unknown") + + # Determine model name and strength + model_name = self.get_filename_from_path(original_path) + strength = None + + if ( + ref.get("strength") is None + and category in {"lora", "loras"} + ): + for node in nodes: + if node_matches_ref(node, ref): + strength = get_lora_model_strength( + node, + widget_index, + ) + break + + if ref.get("strength") is not None: + strength = ref.get("strength") + + model_name, strength = self.adapt_custom_node_loaded_model( + ref, + model_name, + strength, + ) + + # Check if model exists locally + exists = ref.get("exists", False) + + # If URN, resolve to display name + if ref.get("is_urn"): + urn = ref.get("urn", {}) + # Use model name from URN as display name + model_name = ( + f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}" + ) + category = urn.get("type", category) + if category in URN_TYPE_MAP: + category = URN_TYPE_MAP[category] + + loaded_models.append( + { + "name": model_name, + "category": category, + "node_id": node_id, + "widget_index": widget_index, + "node_type": node_type, + "node_title": ref.get("node_title", ""), + "subgraph_id": ref.get("subgraph_id") or "", + "subgraph_name": ref.get("subgraph_name") or "", + "is_top_level": ref.get("is_top_level") is not False, + "locate_node_id": ref.get("locate_node_id"), + "locate_node_type": ref.get("locate_node_type", ""), + "locate_node_title": ref.get("locate_node_title", ""), + "locate_subgraph_id": ( + ref.get("locate_subgraph_id") or "" + ), + "locate_subgraph_name": ( + ref.get("locate_subgraph_name") or "" + ), + "locate_is_top_level": ( + ref.get("locate_is_top_level") is not False + ), + "locate_via_promoted_widget": ref.get( + "locate_via_promoted_widget", False + ), + "exists": exists, + "strength": strength, + "original_path": original_path, + "is_urn": ref.get("is_urn", False), + "custom_node_adapter": ref.get( + "custom_node_adapter" + ), + "active": ref.get("active"), + "connected": ref.get("connected", True), + "resolved_path": ( + path_by_filename.get(model_name) + or path_by_filename.get(original_path) + ), + } + ) + + # Also check node.properties.models for embedded models + total_nodes = len(nodes) + for node_index, node in enumerate(nodes, start=1): + if total_nodes and ( + node_index == total_nodes or node_index % 50 == 0 + ): + update_loaded_progress( + "embedded", + "Checking embedded model metadata...", + percent=interpolate_percent( + 94, 98, node_index, total_nodes + ), + current=node_index, + total=total_nodes, + ) + + node_type = node.get("type", "") + properties = node.get("properties", {}) + if not isinstance(properties, dict): + properties = {} + models_list = properties.get("models", []) + if not isinstance(models_list, list): + models_list = [] + + for model_info in models_list: + if isinstance(model_info, dict): + name = model_info.get("name", "") + directory = model_info.get("directory", "") + node_scope = get_node_scope(node) + + if name: + # Check if this model is already in loaded_models + existing = next( + ( + model + for model in loaded_models + if model.get("original_path") == name + ), + None, + ) + if not existing: + loaded_models.append( + { + "name": self.get_filename_from_path(name), + "category": directory or "checkpoints", + "node_id": node.get("id"), + "widget_index": None, + "node_type": node_type, + "node_title": node.get("title", ""), + "subgraph_id": node_scope["subgraph_id"], + "subgraph_name": node_scope[ + "subgraph_name" + ], + "is_top_level": node_scope[ + "is_top_level" + ], + "exists": True, + "strength": None, + "original_path": name, + "is_urn": False, + } + ) + + update_loaded_progress( + "finalizing", + "Loaded models ready", + percent=99, + current=len(loaded_models), + total=len(loaded_models), + ) + return { + "loaded_models": loaded_models, + "total": len(loaded_models), + } + + update_loaded_progress( + "starting", + "Preparing loaded model scan...", + percent=0, + status="starting", + current=0, + total=0, + ) + + try: + result = await self.asyncio.to_thread(build_loaded_models_response) + except Exception as error: + update_loaded_progress( + "error", + str(error), + percent=100, + status="error", + current=0, + total=0, + ) + self.extension.logger.error( + f"Model Resolver loaded models error: {error}", exc_info=True + ) + return self.web.json_response({"error": str(error)}, status=500) + + update_loaded_progress( + "completed", + "Loaded models ready", + percent=100, + status="completed", + current=result.get("total", 0), + total=result.get("total", 0), + ) + return self.web.json_response(result) + + async def get_loaded_models_progress(self, request): + """Get loaded models inspection progress.""" + return self.get_progress_response( + self.extension.loaded_progress, + request, + param_name="loaded_id", + not_found_payload={ + "status": "unknown", + "stage": "unknown", + "message": "No loaded models progress available", + "percent": 0, + "current": 0, + "total": 0, + }, + ) diff --git a/tests/test_loaded_models_routes.py b/tests/test_loaded_models_routes.py new file mode 100644 index 0000000..01f631e --- /dev/null +++ b/tests/test_loaded_models_routes.py @@ -0,0 +1,428 @@ +import asyncio +import json +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from aiohttp import web + +from core.path_utils import get_filename_from_path +from core.routes.context import RouteContext +from core.routes.loaded_models import register_loaded_model_routes + + +class _Routes: + def __init__(self): + self.handlers = {} + + def get(self, path): + def decorator(handler): + self.handlers[("GET", path)] = handler + return handler + + return decorator + + def post(self, path): + def decorator(handler): + self.handlers[("POST", path)] = handler + return handler + + return decorator + + +def _json_api_endpoint(_name, **_kwargs): + return lambda handler: handler + + +def _build_routes(inventory=None): + routes = _Routes() + extension = SimpleNamespace( + loaded_progress=MagicMock(), + logger=MagicMock(), + _update_loaded_progress=MagicMock(), + _update_workflow_analysis_progress=MagicMock(), + ) + values = { + "adapt_custom_node_loaded_model": MagicMock( + side_effect=lambda _ref, name, strength: (name, strength) + ), + "asyncio": asyncio, + "get_filename_from_path": get_filename_from_path, + "get_progress_response": MagicMock( + return_value={"status": "running"} + ), + "get_workflow_model_inventory": MagicMock( + return_value=inventory + or {"available_models": [], "model_refs": []} + ), + "json_api_endpoint": _json_api_endpoint, + "routes": routes, + "self": extension, + "web": web, + } + register_loaded_model_routes(RouteContext(values)) + return routes.handlers, values + + +def _request(payload=None, loaded_id="loaded-1"): + return SimpleNamespace( + json=AsyncMock(return_value=payload), + match_info={"loaded_id": loaded_id}, + ) + + +def _folder_paths(): + folder_paths = MagicMock() + folder_paths.get_filename_list.return_value = [] + folder_paths.get_full_path.return_value = None + return folder_paths + + +@pytest.mark.asyncio +async def test_loaded_models_route_returns_resolved_local_model(): + inventory = { + "available_models": [ + { + "relative_path": "checkpoints/model.safetensors", + "path": r"C:\models\checkpoints\model.safetensors", + } + ], + "model_refs": [ + { + "original_path": "checkpoints/model.safetensors", + "node_id": 1, + "widget_index": 0, + "node_type": "CheckpointLoaderSimple", + "category": "checkpoints", + "exists": True, + } + ], + } + handlers, values = _build_routes(inventory) + handler = handlers[("POST", "/model_resolver/loaded")] + workflow = {"nodes": [{"id": 1, "type": "CheckpointLoaderSimple"}]} + + with patch.dict(sys.modules, {"folder_paths": _folder_paths()}): + response = await handler(_request({"workflow": workflow})) + + body = json.loads(response.text) + assert response.status == 200 + assert body["total"] == 1 + assert body["loaded_models"][0]["name"] == "model.safetensors" + assert ( + body["loaded_models"][0]["resolved_path"] + == r"C:\models\checkpoints\model.safetensors" + ) + values["get_workflow_model_inventory"].assert_called_once() + values["self"]._update_loaded_progress.assert_any_call( + "", + "completed", + "Loaded models ready", + percent=100, + status="completed", + current=1, + total=1, + ) + + +@pytest.mark.asyncio +async def test_loaded_models_route_includes_urn_and_embedded_models(): + handlers, _ = _build_routes( + {"available_models": [], "model_refs": [ + { + "original_path": "urn:checkpoint:123", + "node_id": 2, + "widget_index": 1, + "node_type": "CustomLoader", + "category": "unknown", + "is_urn": True, + "urn": {"type": "checkpoint", "model_id": 123}, + } + ]} + ) + handler = handlers[("POST", "/model_resolver/loaded")] + workflow = { + "nodes": [ + { + "id": 2, + "type": "CustomLoader", + "properties": { + "models": [ + {"name": "embedded.safetensors", "directory": "loras"} + ] + }, + } + ] + } + + with patch.dict(sys.modules, {"folder_paths": _folder_paths()}): + response = await handler(_request({"workflow": workflow}, "job-2")) + + body = json.loads(response.text) + names = {model["name"] for model in body["loaded_models"]} + assert response.status == 200 + assert "urn:checkpoint:123" in names + assert "embedded.safetensors" in names + assert body["total"] == 2 + + +@pytest.mark.asyncio +async def test_loaded_models_route_returns_error_when_inventory_fails(): + handlers, values = _build_routes() + values["get_workflow_model_inventory"].side_effect = RuntimeError( + "inventory failed" + ) + handler = handlers[("POST", "/model_resolver/loaded")] + + response = await handler( + _request({"workflow": {"nodes": []}, "loaded_id": "failed-job"}) + ) + + assert response.status == 500 + assert json.loads(response.text) == {"error": "inventory failed"} + values["self"].logger.error.assert_called_once() + values["self"]._update_loaded_progress.assert_any_call( + "failed-job", + "error", + "inventory failed", + percent=100, + status="error", + current=0, + total=0, + ) + + +@pytest.mark.asyncio +async def test_loaded_models_progress_route_uses_loaded_id_and_fallback_payload(): + handlers, values = _build_routes() + handler = handlers[("GET", "/model_resolver/loaded-progress/{loaded_id}")] + request = _request(loaded_id="progress-42") + + response = await handler(request) + + assert response == {"status": "running"} + values["get_progress_response"].assert_called_once_with( + values["self"].loaded_progress, + request, + param_name="loaded_id", + not_found_payload={ + "status": "unknown", + "stage": "unknown", + "message": "No loaded models progress available", + "percent": 0, + "current": 0, + "total": 0, + }, + ) + + +@pytest.mark.asyncio +async def test_loaded_models_route_handles_invalid_node_collections_without_folder_paths(): + handlers, _ = _build_routes({"available_models": [], "model_refs": []}) + handler = handlers[("POST", "/model_resolver/loaded")] + workflow = { + "nodes": "not-a-list", + "definitions": { + "subgraphs": [None, {"nodes": "not-a-list"}], + }, + } + + with patch.dict(sys.modules, {"folder_paths": None}): + response = await handler(_request({"workflow": workflow})) + + assert response.status == 200 + assert json.loads(response.text) == {"loaded_models": [], "total": 0} + + +@pytest.mark.asyncio +async def test_loaded_models_route_reads_folder_paths_and_ignores_category_errors(): + folder_paths = MagicMock() + folder_paths.get_filename_list.side_effect = [ + ["folder/model.safetensors"], + RuntimeError("folder list failed"), + [], + [], + [], + ] + folder_paths.get_full_path.return_value = ( + r"C:\models\folder\model.safetensors" + ) + handlers, _ = _build_routes( + { + "available_models": [], + "model_refs": [ + { + "original_path": "folder/model.safetensors", + "node_id": 1, + "category": "checkpoints", + } + ], + } + ) + + with patch.dict(sys.modules, {"folder_paths": folder_paths}): + response = await handlers[("POST", "/model_resolver/loaded")]( + _request({"workflow": {"nodes": []}}) + ) + + body = json.loads(response.text) + assert response.status == 200 + assert body["loaded_models"][0]["resolved_path"] == ( + r"C:\models\folder\model.safetensors" + ) + assert folder_paths.get_filename_list.call_count == 5 + + +@pytest.mark.asyncio +async def test_loaded_models_route_matches_top_level_and_subgraph_lora_refs(): + inventory = { + "available_models": [], + "model_refs": [ + { + "original_path": "top-style.safetensors", + "node_id": 2, + "widget_index": 0, + "node_type": "LoraLoader", + "category": "loras", + }, + { + "original_path": "inner-style.safetensors", + "node_id": 3, + "widget_index": 0, + "node_type": "LoraLoader", + "category": "loras", + "is_top_level": False, + "subgraph_id": "group-1", + "strength": 1.25, + }, + ], + } + handlers, _ = _build_routes(inventory) + workflow = { + "nodes": [ + {"id": 1, "type": "Other"}, + { + "id": 2, + "type": "LoraLoader", + "widgets_values": ["top-style.safetensors", 0.7], + }, + ], + "definitions": { + "subgraphs": [ + None, + {"id": "invalid-group", "nodes": "not-a-list"}, + { + "id": "group-1", + "name": "Group one", + "nodes": [ + { + "id": 3, + "type": "LoraLoader", + "widgets_values": [ + "inner-style.safetensors", + 0.8, + ], + } + ], + }, + ] + }, + } + + with ( + patch.dict(sys.modules, {"folder_paths": None}), + patch( + "core.workflow_analyzer.get_lora_model_strength", + return_value=0.7, + ) as get_strength, + ): + response = await handlers[("POST", "/model_resolver/loaded")]( + _request({"workflow": workflow, "progress_id": "lora-job"}) + ) + + body = json.loads(response.text) + assert response.status == 200 + assert [model["strength"] for model in body["loaded_models"]] == [ + 0.7, + 1.25, + ] + assert body["loaded_models"][1]["subgraph_id"] == "group-1" + assert body["loaded_models"][1]["is_top_level"] is False + assert get_strength.call_count == 1 + + +@pytest.mark.asyncio +async def test_loaded_models_route_uses_adapter_name_and_original_path_fallback(): + inventory = { + "available_models": [ + { + "relative_path": "folder/actual.safetensors", + "path": r"C:\models\actual.safetensors", + } + ], + "model_refs": [ + { + "original_path": "folder/actual.safetensors", + "node_id": 1, + "category": "checkpoints", + } + ], + } + handlers, values = _build_routes(inventory) + values["adapt_custom_node_loaded_model"].side_effect = ( + lambda _ref, _name, _strength: ("display-name.safetensors", 0.5) + ) + + progress_payload = {"processed": 1, "total": 1} + + def inventory_with_progress(_workflow, progress_callback): + progress_callback(progress_payload) + return inventory + + values["get_workflow_model_inventory"].side_effect = inventory_with_progress + + with patch.dict(sys.modules, {"folder_paths": None}): + response = await handlers[("POST", "/model_resolver/loaded")]( + _request({"workflow": {"nodes": [{"id": 1}]}}) + ) + + body = json.loads(response.text) + assert response.status == 200 + assert body["loaded_models"][0]["name"] == "display-name.safetensors" + assert body["loaded_models"][0]["resolved_path"] == ( + r"C:\models\actual.safetensors" + ) + values["self"]._update_workflow_analysis_progress.assert_called_once() + + +@pytest.mark.asyncio +async def test_loaded_models_route_handles_embedded_metadata_edge_cases(): + handlers, _ = _build_routes({"available_models": [], "model_refs": []}) + workflow = { + "nodes": [ + {"id": 1, "type": "NoProperties", "properties": "invalid"}, + { + "id": 2, + "type": "EmbeddedLoader", + "properties": { + "models": [ + "not-a-model", + {"name": ""}, + {"name": "embedded.safetensors"}, + {"name": "embedded.safetensors"}, + ] + }, + }, + ] + } + + with patch.dict(sys.modules, {"folder_paths": None}): + response = await handlers[("POST", "/model_resolver/loaded")]( + _request({"workflow": workflow}) + ) + + body = json.loads(response.text) + assert response.status == 200 + assert body["total"] == 1 + assert body["loaded_models"][0]["name"] == "embedded.safetensors"