refactor(loaded): extract loaded model route handlers into service
This commit is contained in:
Dariusz L
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user