diff --git a/py/manager.py b/py/manager.py index 61a0532..4b8b118 100644 --- a/py/manager.py +++ b/py/manager.py @@ -131,12 +131,11 @@ class ModelManager: basename = os.path.splitext(filename)[0] if is_file else filename extension = os.path.splitext(filename)[1] if is_file else "" - if is_file and extension not in folder_paths.supported_pt_extensions: - return None - - preview_name = utils.get_model_preview_name(entry.path) - preview_ext = f".{preview_name.split('.')[-1]}" - model_preview = f"/model-manager/preview/{folder}/{path_index}/{relative_path.replace(extension, preview_ext)}" + model_preview = None + if is_file: + preview_name = utils.get_model_preview_name(entry.path) + preview_ext = f".{preview_name.split('.')[-1]}" + model_preview = f"/model-manager/preview/{folder}/{path_index}/{relative_path.replace(extension, preview_ext)}" if not os.path.exists(entry.path): utils.print_error(f"{entry.path} is not file or directory.") @@ -151,7 +150,7 @@ class ModelManager: "extension": extension, "pathIndex": path_index, "sizeBytes": stat.st_size if is_file else 0, - "preview": model_preview if is_file else None, + "preview": model_preview, "createdAt": round(stat.st_ctime_ns / 1000000), "updatedAt": round(stat.st_mtime_ns / 1000000), } @@ -160,26 +159,34 @@ class ModelManager: entries: list[os.DirEntry[str]] = [] with os.scandir(directory) as it: for entry in it: - # Skip hidden files - if not include_hidden_files: - if entry.name.startswith("."): - continue - entries.append(entry) - if entry.is_dir(): + if not include_hidden_files and entry.name.startswith("."): + continue + + if entry.is_file(): + extension = os.path.splitext(entry.name)[1] + if extension in folder_paths.supported_pt_extensions: + entries.append(entry) + else: + entries.append(entry) entries.extend(get_all_files_entry(entry.path)) return entries + BATCH_SIZE = 200 + MAX_WORKERS = min(4, os.cpu_count() or 1) + for path_index, base_path in enumerate(folders): if not os.path.exists(base_path): continue file_entries = get_all_files_entry(base_path) - with ThreadPoolExecutor() as executor: - futures = {executor.submit(get_file_info, entry, base_path, path_index): entry for entry in file_entries} - for future in as_completed(futures): - file_info = future.result() - if file_info is None: - continue - result.append(file_info) + + for i in range(0, len(file_entries), BATCH_SIZE): + batch = file_entries[i:i + BATCH_SIZE] + with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: + futures = {executor.submit(get_file_info, entry, base_path, path_index): entry for entry in batch} + for future in as_completed(futures): + file_info = future.result() + if file_info is not None: + result.append(file_info) return result diff --git a/py/utils.py b/py/utils.py index 9f42a2d..c5a8eec 100644 --- a/py/utils.py +++ b/py/utils.py @@ -21,6 +21,9 @@ from . import config VIDEO_EXTENSIONS = ['.mp4', '.webm', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.m4v', '.ogv'] IMAGE_EXTENSIONS = ['.webp', '.png', '.jpg', '.jpeg', '.gif', '.bmp'] +# Preview extensions in priority order (videos first, then images) +PREVIEW_EXTENSIONS = ['.webm', '.mp4', '.webp', '.png', '.jpg', '.jpeg', '.gif', '.bmp'] + # Content type mappings VIDEO_CONTENT_TYPE_MAP = { 'video/mp4': '.mp4', @@ -268,40 +271,52 @@ def get_model_metadata(filename: str): return {} -def get_model_all_previews(model_path: str): +def _check_preview_variants(base_dirname: str, basename: str, extensions: list[str]) -> list[str]: + """Check for preview files with given extensions and return found files""" + found = [] + for ext in extensions: + # Direct match (basename.ext) + preview_file = f"{basename}{ext}" + if os.path.isfile(join_path(base_dirname, preview_file)): + found.append(preview_file) + + # Preview variant (basename.preview.ext) + preview_file = f"{basename}.preview{ext}" + if os.path.isfile(join_path(base_dirname, preview_file)): + found.append(preview_file) + return found + + +def _get_preview_path(model_path: str, extension: str) -> str: + """Generate preview file path with given extension""" + basename = os.path.splitext(model_path)[0] + return f"{basename}{extension}" + + +def get_model_all_previews(model_path: str) -> list[str]: + """Get all preview files for a model""" base_dirname = os.path.dirname(model_path) - files = search_files(base_dirname) - files = folder_paths.filter_files_content_types(files, ["video", "image"]) - basename = os.path.splitext(os.path.basename(model_path))[0] - output: list[str] = [] - for file in files: - file_basename = os.path.splitext(file)[0] - if file_basename == basename: - output.append(file) - if file_basename == f"{basename}.preview": - output.append(file) - return output + return _check_preview_variants(base_dirname, basename, PREVIEW_EXTENSIONS) -def get_model_preview_name(model_path: str): - """ - Get the preview file name for a model. Checks for images and videos in all supported formats. - Returns the first available preview file or 'no-preview.png' if none found. - """ +def get_model_preview_name(model_path: str) -> str: + """Get the first available preview file or 'no-preview.png' if none found""" base_dirname = os.path.dirname(model_path) basename = os.path.splitext(os.path.basename(model_path))[0] - # Prefer previews with these extensions in this order - preview_extensions = ['.webm', '.mp4', '.webp'] - for ext in preview_extensions: + for ext in PREVIEW_EXTENSIONS: + # Check direct match first preview_name = f"{basename}{ext}" - if os.path.exists(join_path(base_dirname, preview_name)): + if os.path.isfile(join_path(base_dirname, preview_name)): + return preview_name + + # Check preview variant + preview_name = f"{basename}.preview{ext}" + if os.path.isfile(join_path(base_dirname, preview_name)): return preview_name - # Fallback to any available preview files - all_previews = get_model_all_previews(model_path) - return all_previews[0] if len(all_previews) > 0 else "no-preview.png" + return "no-preview.png" from PIL import Image @@ -309,34 +324,19 @@ from io import BytesIO def remove_model_preview(model_path: str): - """ - Remove preview files for a model. - """ - basename = os.path.splitext(model_path)[0] + """Remove all preview files for a model""" base_dirname = os.path.dirname(model_path) + basename = os.path.splitext(os.path.basename(model_path))[0] - # Remove all preview files - for ext in VIDEO_EXTENSIONS + IMAGE_EXTENSIONS: - preview_path = f"{basename}{ext}" + previews = _check_preview_variants(base_dirname, basename, PREVIEW_EXTENSIONS) + for preview in previews: + preview_path = join_path(base_dirname, preview) if os.path.exists(preview_path): os.remove(preview_path) - - # Also check for .preview variants - files = search_files(base_dirname) - model_name = os.path.splitext(os.path.basename(model_path))[0] - for file in files: - if file.startswith(f"{model_name}.preview"): - file_path = join_path(base_dirname, file) - if os.path.exists(file_path): - os.remove(file_path) def save_model_preview(model_path: str, file_or_url: Any, platform: Optional[str] = None): - """ - Save a preview file for a model. - Images are converted to WebP, videos are saved in their original format. - """ - basename = os.path.splitext(model_path)[0] + """Save a preview file for a model. Images -> WebP, videos -> original format""" # Download file if it is a URL if type(file_or_url) is str: @@ -357,15 +357,13 @@ def save_model_preview(model_path: str, file_or_url: Any, platform: Optional[str if content_type.startswith("video/"): # Save video in original format # Try to get extension from URL or content-type - ext = _get_video_extension_from_url(url) or _get_extension_from_content_type(content_type) - if not ext: - ext = '.mp4' # Default fallback - preview_path = f"{basename}{ext}" + ext = _get_video_extension_from_url(url) or _get_extension_from_content_type(content_type) or '.mp4' + preview_path = _get_preview_path(model_path, ext) with open(preview_path, 'wb') as f: f.write(content) else: # Default to image processing for unknown or image types - preview_path = f"{basename}.webp" + preview_path = _get_preview_path(model_path, ".webp") image = Image.open(BytesIO(content)) image.save(preview_path, "WEBP") @@ -384,17 +382,15 @@ def save_model_preview(model_path: str, file_or_url: Any, platform: Optional[str if content_type.startswith("video/"): # Save video in original format for now, consider transcoding to webm to follow the pattern for images converting to webp - ext = os.path.splitext(filename.lower())[1] - if not ext: - ext = '.mp4' # Default fallback - preview_path = f"{basename}{ext}" + ext = os.path.splitext(filename.lower())[1] or '.mp4' + preview_path = _get_preview_path(model_path, ext) file_obj.file.seek(0) content = file_obj.file.read() with open(preview_path, 'wb') as f: f.write(content) elif content_type.startswith("image/"): # Convert image to webp - preview_path = f"{basename}.webp" + preview_path = _get_preview_path(model_path, ".webp") image = Image.open(file_obj.file) image.save(preview_path, "WEBP") else: diff --git a/pyproject.toml b/pyproject.toml index 47b2b21..e7469ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-model-manager" description = "Manage models: browsing, download and delete." -version = "2.8.1" +version = "2.8.2" license = { file = "LICENSE" } dependencies = ["markdownify"]