Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f3de2006ef | ||
|
|
0295dd6288 | ||
|
|
4f9a437725 | ||
|
|
815a483cf0 | ||
|
|
ae37765017 | ||
|
|
ebef300279 |
@@ -131,12 +131,15 @@ 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
|
||||
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)}"
|
||||
|
||||
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.")
|
||||
return None
|
||||
|
||||
stat = entry.stat()
|
||||
return {
|
||||
@@ -147,35 +150,45 @@ 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),
|
||||
}
|
||||
|
||||
def get_all_files_entry(directory: str):
|
||||
entries: list[os.DirEntry[str]] = []
|
||||
if not os.path.exists(directory):
|
||||
return []
|
||||
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
|
||||
|
||||
|
||||
108
py/utils.py
108
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',
|
||||
@@ -43,7 +46,7 @@ def print_warning(msg, *args, **kwargs):
|
||||
|
||||
|
||||
def print_error(msg, *args, **kwargs):
|
||||
logging.error(f"[{config.extension_tag}] {msg}", *args, **kwargs)
|
||||
logging.error(f"[{config.extension_tag}][ERROR] {msg}", *args, **kwargs)
|
||||
logging.debug(traceback.format_exc())
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-model-manager"
|
||||
description = "Manage models: browsing, download and delete."
|
||||
version = "2.8.0"
|
||||
version = "2.8.3"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = ["markdownify"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user