"""Model folder discovery, listing, search, and presentation lookup.""" import asyncio import os import urllib.parse from aiohttp import web import folder_paths from .metadata import get_metadata try: from ..model_sidecars import HIDDEN_DIR except ImportError: from model_sidecars import HIDDEN_DIR from .model_constants import MEDIA_EXTENSIONS, PREVIEW_SUFFIXES, is_model_file from .model_media import _cache_token, _preview_url_for_model from .folder_types import ( get_active_folder_types, get_active_model_roots, get_active_physical_basenames, get_folder_view_mode, ) from .path_utils import model_reference_within, resolve_folder_subdir def _folder_tree(base_dir): """{'/'-relative folder: {path, name, has_models, model_count}} under a model root; the .amb folders hold side files only (model_sidecars).""" tree = {} for root, dirs, files in os.walk(base_dir): dirs[:] = [name for name in dirs if name != HIDDEN_DIR] count = sum(1 for f in files if is_model_file(f)) rel = os.path.relpath(root, base_dir) rel = '/' if rel == '.' else '/' + rel.replace('\\', '/') tree[rel] = {"path": rel, "name": os.path.basename(root) if rel != '/' else '[Root]', "has_models": count > 0, "model_count": count} return tree def _collect_folder_models(target_dir, folder_type, path_idx, rel_subfolder, page, limit): """Collect one folder in a worker thread with a single directory listing.""" try: with os.scandir(target_dir) as iterator: file_entries = {} for entry in iterator: try: if entry.is_file(): file_entries[entry.name] = entry except OSError: continue except OSError: return {"models": [], "total": 0, "page": page, "limit": limit} valid_files = sorted( (name for name in file_entries if is_model_file(name)), key=str.lower, ) total = len(valid_files) if limit > 0: start = max(0, (page - 1) * limit) sliced = valid_files[start:start + limit] else: sliced = valid_files q_type = urllib.parse.quote(folder_type) q_idx = str(path_idx) q_sub = urllib.parse.quote(rel_subfolder.strip('/')) if rel_subfolder and rel_subfolder != '/' else "" models = [] for filename in sliced: file_path = os.path.join(target_dir, filename) metadata = get_metadata(file_path) base_name = os.path.splitext(filename)[0] preview_file = next( (base_name + suffix for suffix in PREVIEW_SUFFIXES + MEDIA_EXTENSIONS if base_name + suffix in file_entries), None, ) preview_url = "" if preview_file: q_file = urllib.parse.quote(preview_file) try: preview_version = file_entries[preview_file].stat().st_mtime_ns except OSError: preview_version = 0 preview_url = ( f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}" f"&filename={q_file}&t={preview_version}" ) try: model_stat = file_entries[filename].stat() size_bytes = model_stat.st_size size_mb = round(size_bytes / (1024 * 1024), 2) except OSError: size_mb = 0 size_bytes = 0 models.append({ "filename": filename, "file_path": os.path.abspath(file_path), "size_mb": size_mb, "size_bytes": size_bytes, "metadata": metadata, "preview_url": preview_url, "type": folder_type, "path_idx": path_idx, "subfolder": rel_subfolder, }) return {"models": models, "total": total, "page": page, "limit": limit} def _collect_folders(): mode = get_folder_view_mode() result = [] seen_dirs = set() if mode == "physical": active_bns = get_active_physical_basenames() all_paths_info = [] for t in folder_paths.folder_names_and_paths.keys(): try: paths = folder_paths.get_folder_paths(t) if not paths: continue for path_idx, base_dir in enumerate(paths): if not os.path.exists(base_dir): continue real_path = os.path.realpath(base_dir) if real_path in seen_dirs: continue seen_dirs.add(real_path) bn = os.path.basename(os.path.normpath(base_dir)) all_paths_info.append({ "t": t, "path_idx": path_idx, "base_dir": base_dir, "bn": bn }) except: pass for target_bn in active_bns: matched = [p for p in all_paths_info if p["bn"] == target_bn] if not matched: continue for item in matched: base_dir = item["base_dir"] tree = _folder_tree(base_dir) label = item["bn"] if len(matched) > 1: label += f" ({item['path_idx'] + 1})" result.append({ "type": item["t"], "path_idx": item["path_idx"], "label": label, "folders": tree }) else: types = get_active_folder_types() for t in types: try: paths = folder_paths.get_folder_paths(t) except Exception: continue if not paths: continue for path_idx, base_dir in enumerate(paths): if not os.path.exists(base_dir): continue real_path = os.path.realpath(base_dir) if real_path in seen_dirs: continue seen_dirs.add(real_path) tree = _folder_tree(base_dir) try: folder_basename = os.path.basename(os.path.normpath(base_dir)) if not folder_basename: folder_basename = t except: folder_basename = t label = folder_basename basenames = [] try: basenames = [os.path.basename(os.path.normpath(p)) for p in paths] except: pass if basenames.count(folder_basename) > 1: label += f" ({path_idx + 1})" result.append({ "type": t, "path_idx": path_idx, "label": label, "folders": tree }) return {"folders": result} async def api_get_folders(request): return web.json_response(await asyncio.to_thread(_collect_folders)) async def api_get_models(request): folder_type = request.query.get('type', 'checkpoints') subfolder = request.query.get('subfolder', '/') page = int(request.query.get('page', 1)) limit = int(request.query.get('limit', 0)) try: path_idx = int(request.query.get('path_idx', 0)) except: path_idx = 0 try: paths = folder_paths.get_folder_paths(folder_type) except Exception: return web.json_response({"models": [], "total": 0}) try: base_dir, target_dir = resolve_folder_subdir(folder_type, path_idx, subfolder) except (ValueError, KeyError): return web.Response(status=400, text='Invalid subfolder') rel_subfolder = "" if subfolder == '/' else subfolder.strip('/\\') if not os.path.exists(target_dir): return web.json_response({"models": [], "total": 0}) payload = await asyncio.to_thread( _collect_folder_models, target_dir, folder_type, path_idx, rel_subfolder, page, limit, ) return web.json_response(payload) def _iter_search_models(): for folder_type in folder_paths.folder_names_and_paths.keys(): try: paths = folder_paths.get_folder_paths(folder_type) except Exception: continue for path_idx, base_dir in enumerate(paths): if not os.path.isdir(base_dir): continue for root, _, files in os.walk(base_dir): for filename in files: if is_model_file(filename): yield folder_type, path_idx, base_dir, os.path.join(root, filename) def _find_model_sync(search): normalized_search = search.replace('\\', '/').lower() for folder_type, path_idx, base_dir, file_path in _iter_search_models(): normalized_path = file_path.replace(os.sep, '/').lower() if normalized_search in os.path.basename(file_path).lower() or normalized_search in normalized_path: return _model_info_for_path(folder_type, path_idx, base_dir, file_path) for folder_type, path_idx, base_dir, file_path in _iter_search_models(): metadata = get_metadata(file_path) searchable_names = (metadata.get('custom_name', ''), metadata.get('name', '')) if any(normalized_search in str(name).lower() for name in searchable_names if name): return _model_info_for_path(folder_type, path_idx, base_dir, file_path) return None async def api_find_model(request): search = request.query.get('search', '').strip() if not search: return web.json_response({"status": "error", "message": "No search query provided"}) result = await asyncio.to_thread(_find_model_sync, search) if not result: return web.json_response({"status": "error", "message": "Model not found"}) return web.json_response({ "status": "success", "model": result, "type": result["type"], "path_idx": result["path_idx"], "subfolder": result["subfolder"], }) # Prompt Notes ask for this base model to list models whose metadata names none. UNLABELED_BASE_MODEL = "__unlabeled__" def _normalised_base(value): return str(value or "").strip().lower().replace(" ", "") def _walk_model_files(target_types): """Yield (folder_type, path_idx, base_dir, root, filename, metadata) once per real file.""" seen_files = set() for folder_type in target_types: try: paths = folder_paths.get_folder_paths(folder_type) except Exception: continue for path_idx, base_dir in enumerate(paths or []): if not os.path.exists(base_dir): continue for root, _, files in os.walk(base_dir): for filename in files: if not is_model_file(filename): continue file_path = os.path.join(root, filename) real_path = os.path.realpath(file_path) if real_path in seen_files: continue seen_files.add(real_path) yield folder_type, path_idx, base_dir, root, filename, get_metadata(file_path) def _base_matches(requested, metadata): model_base = _normalised_base(metadata.get("baseModel")) if requested == UNLABELED_BASE_MODEL: return not model_base wanted = _normalised_base(requested) return bool(wanted and model_base and (wanted in model_base or model_base in wanted)) def _compatible_models_sync(base_model, target_types): compatible_models = [] for folder_type, path_idx, base_dir, root, filename, meta in _walk_model_files(target_types): if not _base_matches(base_model, meta): continue rel_subfolder = os.path.relpath(root, base_dir) rel_subfolder = '/' if rel_subfolder == '.' else '/' + rel_subfolder.replace('\\', '/') base_name = os.path.splitext(filename)[0] preview_file = None for ext in PREVIEW_SUFFIXES + MEDIA_EXTENSIONS: if os.path.exists(os.path.join(root, base_name + ext)): preview_file = base_name + ext break preview_url = "" if preview_file: q_sub = urllib.parse.quote(rel_subfolder.strip('/')) if rel_subfolder != '/' else "" preview_url = ( f"/anomalous/image?type={urllib.parse.quote(folder_type)}&path_idx={path_idx}" f"&subfolder={q_sub}&filename={urllib.parse.quote(preview_file)}" ) try: size_bytes = os.path.getsize(os.path.join(root, filename)) except OSError: size_bytes = 0 compatible_models.append({ "type": folder_type, "path_idx": path_idx, "subfolder": rel_subfolder, "filename": filename, "size_mb": round(size_bytes / (1024 * 1024), 1), "size_bytes": size_bytes, "preview_url": preview_url, "metadata": meta, }) return compatible_models async def api_compatible_models(request): base_model = request.query.get('base_model', '') if not base_model: return web.json_response({"models": []}) target_types = [t.strip() for t in request.query.get('target_type', 'loras').split(',') if t.strip()] models = await asyncio.to_thread(_compatible_models_sync, base_model, target_types) return web.json_response({"models": models}) def _base_models_sync(target_types): base_models = set() unlabeled = False for *_, meta in _walk_model_files(target_types): label = str(meta.get("baseModel") or "").strip() if label: base_models.add(label) else: unlabeled = True return sorted(base_models), unlabeled async def api_base_models(request): base_models, unlabeled = await asyncio.to_thread(_base_models_sync, get_active_folder_types()) return web.json_response({"base_models": base_models, "unlabeled": unlabeled}) def _model_info_for_path(folder_type, path_idx, base_dir, file_path): root = os.path.dirname(file_path) rel_subfolder = os.path.relpath(root, base_dir) if rel_subfolder == '.': rel_subfolder = '/' else: rel_subfolder = '/' + rel_subfolder.replace(os.sep, '/') try: size_bytes = os.path.getsize(file_path) except OSError: size_bytes = 0 return { "type": folder_type, "path_idx": path_idx, "subfolder": rel_subfolder, "filename": os.path.basename(file_path), "file_path": os.path.abspath(file_path), "size_bytes": size_bytes, "size_mb": round(size_bytes / (1024 * 1024), 2), "preview_url": _preview_url_for_model(folder_type, path_idx, base_dir, file_path), "metadata": get_metadata(file_path), } def _allowed_folder_types(requested_types=None): available = list(folder_paths.folder_names_and_paths.keys()) if requested_types is None: return available if not isinstance(requested_types, list): return [] allowed = set(available) return list(dict.fromkeys( folder_type for folder_type in requested_types if isinstance(folder_type, str) and folder_type in allowed )) def _resolve_paths_to_model_info_sync(paths, folder_types=None, exact_only=False): requested = [ (path, path.replace('\\', '/').lower(), path.replace('\\', '/')) for path in paths if isinstance(path, str) ] exact_results = {} roots = [] for folder_type in _allowed_folder_types(folder_types): try: folder_dirs = folder_paths.get_folder_paths(folder_type) except Exception: continue for path_idx, base_dir in enumerate(folder_dirs or []): if not os.path.isdir(base_dir): continue real_base_dir = os.path.realpath(base_dir) roots.append((folder_type, path_idx, real_base_dir)) for original, _, relative_path in requested: try: candidate = model_reference_within(real_base_dir, relative_path.replace('/', os.sep)) except ValueError: continue if os.path.isfile(candidate) and is_model_file(candidate): # The candidate lies under the real folder: a folder that is a link (a junction # to another drive, a short 8.3 name) would give a subfolder of "../.." or fail. exact_results[original] = _model_info_for_path(folder_type, path_idx, real_base_dir, candidate) unresolved = [(original, normalized) for original, normalized, _ in requested if original not in exact_results] if exact_only or not unresolved: return exact_results wanted_relpaths = {normalized for _, normalized in unresolved} wanted_basenames = {normalized.rsplit('/', 1)[-1] for _, normalized in unresolved} rel_matches = {} basename_matches = {} for folder_type, path_idx, base_dir in roots: for root, _, files in os.walk(base_dir): for filename in files: if not is_model_file(filename): continue rel_path = os.path.relpath(os.path.join(root, filename), base_dir).replace(os.sep, '/').lower() basename = filename.lower() if rel_path not in wanted_relpaths and basename not in wanted_basenames: continue candidate_info = _model_info_for_path( folder_type, path_idx, base_dir, os.path.join(root, filename), ) if rel_path in wanted_relpaths: rel_matches[rel_path] = candidate_info if basename in wanted_basenames: basename_matches[basename] = candidate_info model_info = dict(exact_results) for original, normalized in unresolved: if normalized in rel_matches: model_info[original] = rel_matches[normalized] else: basename = normalized.rsplit('/', 1)[-1] if basename in basename_matches: model_info[original] = basename_matches[basename] return model_info def _resolve_paths_to_previews_sync(paths, folder_types=None, exact_only=False): model_info = _resolve_paths_to_model_info_sync(paths, folder_types, exact_only) return {path: item.get("preview_url", "") for path, item in model_info.items()} async def api_resolve_paths_to_previews(request): try: data = await request.json() paths = data.get('paths', []) folder_types = data.get('folder_types') context_requests = data.get('context_requests', []) exact_only = data.get('exact_only') is True except: return web.json_response({"previews": {}, "models": {}, "context_models": {}}) model_info = await asyncio.to_thread(_resolve_paths_to_model_info_sync, paths, folder_types, exact_only) context_models = {} if isinstance(context_requests, list): for item in context_requests[:16]: if not isinstance(item, dict): continue context_path = item.get('path') if not isinstance(context_path, str) or not context_path: continue resolved = await asyncio.to_thread( _resolve_paths_to_model_info_sync, [context_path], item.get('folder_types'), item.get('exact_only') is True or exact_only, ) if context_path in resolved: context_key = item.get('key') if not isinstance(context_key, str) or not context_key: context_key = context_path context_models[context_key] = resolved[context_path] previews = {path: item.get("preview_url", "") for path, item in model_info.items()} return web.json_response({ "previews": previews, "models": model_info, "context_models": context_models, }) def _collect_all_scan_models(page, limit): all_tuples = [] for model_root in get_active_model_roots(): t = model_root["type"] path_idx = model_root["path_idx"] base_dir = model_root["base_dir"] for root, dirs, files in os.walk(base_dir): for f in files: if is_model_file(f): all_tuples.append((t, path_idx, root, base_dir, f)) all_tuples.sort(key=lambda x: (x[0], x[4].lower())) total = len(all_tuples) if limit > 0: start = (page - 1) * limit end = start + limit sliced = all_tuples[start:end] else: sliced = all_tuples all_models = [] for t, path_idx, root, base_dir, f in sliced: file_path = os.path.join(root, f) rel_subfolder = os.path.relpath(root, base_dir) if rel_subfolder == '.': rel_subfolder = '' else: rel_subfolder = rel_subfolder.replace('\\', '/') try: size_bytes = os.path.getsize(file_path) size_mb = round(size_bytes / (1024 * 1024), 2) except: size_bytes = 0; size_mb = 0 meta = get_metadata(file_path) base_name = os.path.splitext(f)[0] preview_file = None for ext in PREVIEW_SUFFIXES + MEDIA_EXTENSIONS: if os.path.exists(os.path.join(root, base_name + ext)): preview_file = base_name + ext break preview_url = "" if preview_file: q_type = urllib.parse.quote(t) q_idx = str(path_idx) q_sub = urllib.parse.quote(rel_subfolder.strip('/')) if rel_subfolder and rel_subfolder != '/' else "" q_file = urllib.parse.quote(preview_file) mtime = _cache_token(os.path.join(root, preview_file)) preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}&t={mtime}" all_models.append({ "type": t, "path_idx": path_idx, "subfolder": rel_subfolder, "filename": f, "size_mb": size_mb, "size_bytes": size_bytes, "preview_url": preview_url, "metadata": meta }) return {"models": all_models, "total": total, "page": page, "limit": limit} async def api_get_all_scan_models(request): page = int(request.query.get('page', 1)) limit = int(request.query.get('limit', 0)) payload = await asyncio.to_thread(_collect_all_scan_models, page, limit) return web.json_response(payload)