"""Where a model loads, for the AI: search_models and get_model (mcp_tools.py finds the files) say which node types list each file among their choices (ComfyUI's /object_info, kept a little while by mcp_canvas_edit.py), and, when a file's header says it is another kind than the folder it is in (model_placement.py, the models page's Tidy view), where it belongs and that nothing loads it where it is. So the AI picks the right node at once, or has the model moved, instead of searching ComfyUI itself. move_model moves one model with its covers and notes as Tidy does, behind the "files" gate (the user says yes first unless they let it go ahead).""" import os import threading import folder_paths from . import mcp_canvas_edit, mcp_fit, mcp_tools, model_advice, model_placement from .mcp_actions import _action from .mcp_tools import _LIMIT, _QUERY, ToolError from .model_constants import MODEL_EXTENSIONS from .model_import import Refused, import_types MAX_LOADERS = 6 SEARCH_WORDS = 4 # trigger words a found model shows; get_model has them all _lock = threading.Lock() _index = {"of": None, "inputs": [], "by_choice": {}} def _plain(name): return str(name).replace("\\", "/") def _choices(conf): if not isinstance(conf, list) or not conf: return None if isinstance(conf[0], list): return conf[0] options = conf[1] if len(conf) > 1 and isinstance(conf[1], dict) else {} return options.get("options") if conf[0] == "COMBO" and isinstance(options.get("options"), list) else None def _built(info): """The drop-downs of model files in node definitions `info`: [(node type, choices)] and choice -> their places.""" with _lock: if _index["of"] is info: return _index inputs, by_choice = [], {} for node_type, item in info.items(): for section in ("required", "optional"): for conf in ((item.get("input") or {}).get(section) or {}).values(): listed = [_plain(choice) for choice in _choices(conf) or [] if isinstance(choice, str)] if not any(choice.lower().endswith(MODEL_EXTENSIONS) for choice in listed): continue for choice in listed: by_choice.setdefault(choice, []).append(len(inputs)) inputs.append((node_type, frozenset(listed))) _index.update(of=info, inputs=inputs, by_choice=by_choice) return _index def _folder(folder_type): return frozenset(_plain(name) for name in folder_paths.get_filename_list(folder_type)) def _readers(ctx, folder_type, rel=None): """Node types with a drop-down of the whole folder `folder_type` (one that lists `rel`, when given).""" index = _built(mcp_canvas_edit.object_info(ctx)) folder = _folder(folder_type) if not folder: return [] places = index["by_choice"].get(_plain(rel), []) if rel else range(len(index["inputs"])) found = [] for at in places: node_type, choices = index["inputs"][at] if folder <= choices and node_type not in found: found.append(node_type) return found[:MAX_LOADERS] def _about(ctx, model): """`model` (from mcp_tools) with what loads it: loaded_by, or belongs_in and placement when it is misfiled.""" folder_type, root, base_dir, path = mcp_tools._resolve_model(model["model_id"]) rel = os.path.relpath(path, base_dir).replace(os.sep, "/") home = model_placement.belongs(model_placement.one(folder_type, root, rel, path)) if not home or home["works"]: model["loaded_by"] = _readers(ctx, folder_type, rel) return model kind, to = home["kind"], home["to"] readers = _readers(ctx, kind) words = (f"Its header says it is a {kind} file, but it is in {folder_type}: the nodes listing it there cannot load it" + (f"; {', '.join(readers)} load {kind} files." if readers else ".")) there = os.path.join(folder_paths.get_folder_paths(kind)[to["root"]], *to["rel"].split("/")) if os.path.isfile(there): words += (f" A file of the same name is already in {kind} (model_id {kind}:{to['root']}:{to['rel']}): use that one. " "If it is the same file, this copy is extra; the user can recycle it in Anomalous's models page (Tidy).") else: words += f" move_model with to_type {kind} moves it there (the user confirms first)." model.update(belongs_in=kind, placement=words) return model def _compact(found): """search_models' answer small enough to be read (mcp_fit.BUDGET): what a model_id says already (file, folder, type) left out, loaded_by once for a type when its models share it, a few trigger words, and as many models as fit, the rest by name (more_names), so one answer shows all there is and the AI need not search again to see the rest.""" models, kinds = found["models"], {} for model in models: kinds.setdefault(model["type"], []).append(model) shared = {kind: group[0]["loaded_by"] for kind, group in kinds.items() if all("loaded_by" in model and model["loaded_by"] == group[0]["loaded_by"] for model in group)} for model in models: for key in ("file", "folder"): model.pop(key, None) if model.pop("type") in shared: model.pop("loaded_by") words = model.get("trigger_words") or [] if len(words) > SEARCH_WORDS: model.update(trigger_words=words[:SEARCH_WORDS], more_trigger_words=len(words) - SEARCH_WORDS) answer = {"total_matches": found["total_matches"], **({"loaded_by": shared} if shared else {}), "models": models} everything = list(models) while mcp_fit.size(answer) > mcp_fit.BUDGET and len(models) > 1: models.pop() answer["more_names"] = [str(model.get("name") or model["model_id"])[:mcp_fit.NAME] for model in everything[len(models):]] answer["left_out"] = (f"{len(answer['more_names'])} more, named in more_names to keep this answer readable: " "get_model or a query with the name gives one in full") return answer def search_models(ctx, query="", type="", base_model="", source="", limit=20): found = mcp_tools.search_models(query, type, base_model, source, limit) for model in found["models"]: _about(ctx, model) if model["type"] == "loras": # the weight its author's pictures use: picking one needs no get_model each path = mcp_tools._resolve_model(model["model_id"])[3] weight = model_advice.author_examples(path, "loras", (mcp_tools.get_metadata(path) or {}).get("hash", "")).get("lora_weight") if weight: model["author_weight"] = weight return _compact(found) def get_model(ctx, model_id, include_cover=False): detail, extra = mcp_tools.get_model(model_id, include_cover) folder_type, _idx, _base, path = mcp_tools._resolve_model(model_id) examples = model_advice.author_examples(path, folder_type, detail.get("sha256", "")) if examples: detail["author_examples"] = examples return _about(ctx, detail), extra def move_model(ctx, model_id, to_type): folder_type, root, base_dir, path = mcp_tools._resolve_model(model_id) rel = os.path.relpath(path, base_dir).replace(os.sep, "/") item = model_placement.one(folder_type, root, rel, path) home = model_placement.belongs(item) to_type = str(to_type or "") if to_type == folder_type: raise ToolError(f"It is in {to_type} already.") if to_type not in import_types(): raise ToolError(f"to_type is one of {', '.join(import_types())}.") to = home["to"] if home and home["kind"] == to_type else model_placement.home_for(item, to_type) if not to: raise ToolError(f"This ComfyUI has no {to_type} folder.") if os.path.exists(os.path.join(folder_paths.get_folder_paths(to_type)[to["root"]], *to["rel"].split("/"))): raise ToolError(f"A file of the same name is already in {to_type} (model_id {to_type}:{to['root']}:{to['rel']}); " "nothing was moved. Use that one; the user can recycle an extra copy in Anomalous's models page (Tidy).") try: moved = model_placement.move({"type": folder_type, "root": root, "rel": rel, "to_type": to_type, "to_root": to["root"], "to_rel": to["rel"]}) except Refused as error: raise ToolError(f"Not moved ({error.code}).") mcp_canvas_edit.forget_object_info() answer = {"moved": {"from": f"{folder_type}/{rel}", "to": f"{moved['type']}/{moved['rel']}"}, "model_id": f"{moved['type']}:{moved['root']}:{moved['rel']}", "loaded_by": _readers(ctx, moved["type"], moved["rel"])} if moved["left"]: answer["left_behind"] = moved["left"] # covers or notes that could not go along return answer MODEL_TOOLS = dict([ _action("search_models", "Search models", "Model files by name, trigger word or base model: model_id, and loaded_by (the node types that list " "the file). A misfiled one comes with belongs_in: nothing loads it where it is.", search_models, { "query": _QUERY, "type": {"type": "string", "description": "A models folder: checkpoints, loras, vae…"}, "base_model": {"type": "string", "description": "e.g. SDXL, Pony, Flux."}, "source": {"type": "string", "enum": ["civitai", "inferred_from_file", "not_scanned"], "description": "Where the model's information came from."}, "limit": _LIMIT}, read_only=True), _action("get_model", "Model details", "Everything known about one model: Civitai name, version and link, trigger words, base model, description, " "the user's own name and notes, author_examples (the sampler, steps, CFG, size and, for a LoRA, the weight " "its author's example pictures used: a good start), and loaded_by (or belongs_in and placement) as in " "search_models. include_cover attaches its cover picture.", get_model, {"model_id": {"type": "string", "description": "From search_models."}, "include_cover": {"type": "boolean", "default": False}}, ["model_id"], read_only=True), _action("move_model", "Move a model to its folder", "Moves a model the user has in the wrong models folder (search_models says belongs_in) to where it loads, " "with its covers and notes, as Anomalous's Tidy view does; never over a file. The answer gives its new " "model_id and loaded_by; the canvas lists it from then on.", move_model, {"model_id": {"type": "string", "description": "From search_models."}, "to_type": {"type": "string", "description": "The models folder to move it to: belongs_in from search_models."}}, ["model_id", "to_type"], destructive=True), ])