TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add support for video previews (#197)

Browse Source 
* add support for video previews

* fix two cases where video previews did not show
This commit is contained in:
Kevin Lewis
2025-08-14 22:12:13 -04:00
committed by GitHub
parent c96a164f68
commit 71a200ed5c
13 changed files with 267 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

183
py/utils.py
View File

@@ -17,6 +17,22 @@ from aiohttp import web
from typing import Any, Optional
from . import config
# Media file extensions
VIDEO_EXTENSIONS = ['.mp4', '.webm', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.m4v', '.ogv']
IMAGE_EXTENSIONS = ['.webp', '.png', '.jpg', '.jpeg', '.gif', '.bmp']
# Content type mappings
VIDEO_CONTENT_TYPE_MAP = {
'video/mp4': '.mp4',
'video/webm': '.webm',
'video/quicktime': '.mov',
'video/x-msvideo': '.avi',
'video/x-matroska': '.mkv',
'video/x-flv': '.flv',
'video/x-ms-wmv': '.wmv',
'video/ogg': '.ogv',
}
def print_info(msg, *args, **kwargs):
logging.info(f"[{config.extension_tag}] {msg}", *args, **kwargs)
@@ -252,10 +268,10 @@ def get_model_metadata(filename: str):
return {}
def get_model_all_images(model_path: str):
def get_model_all_previews(model_path: str):
base_dirname = os.path.dirname(model_path)
files = search_files(base_dirname)
files = folder_paths.filter_files_content_types(files, ["image"])
files = folder_paths.filter_files_content_types(files, ["video", "image"])
basename = os.path.splitext(os.path.basename(model_path))[0]
output: list[str] = []
@@ -269,78 +285,135 @@ def get_model_all_images(model_path: str):
def get_model_preview_name(model_path: str):
images = get_model_all_images(model_path)
basename = os.path.splitext(os.path.basename(model_path))[0]
for image in images:
image_name = os.path.splitext(image)[0]
image_ext = os.path.splitext(image)[1]
if image_name == basename and image_ext.lower() == ".webp":
return image
return images[0] if len(images) > 0 else "no-preview.png"
def get_model_all_videos(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.
"""
base_dirname = os.path.dirname(model_path)
files = search_files(base_dirname)
files = folder_paths.filter_files_content_types(files, ["video"])
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
# Prefer previews with these extensions in this order
preview_extensions = ['.webm', '.mp4', '.webp']
for ext in preview_extensions:
preview_name = f"{basename}{ext}"
if os.path.exists(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"
from PIL import Image
from io import BytesIO
def remove_model_preview_image(model_path: str):
def remove_model_preview(model_path: str):
"""
Remove preview files for a model.
"""
basename = os.path.splitext(model_path)[0]
preview_path = f"{basename}.webp"
if os.path.exists(preview_path):
os.remove(preview_path)
base_dirname = os.path.dirname(model_path)
# Remove all preview files
for ext in VIDEO_EXTENSIONS + IMAGE_EXTENSIONS:
preview_path = f"{basename}{ext}"
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_image(model_path: str, image_file_or_url: Any, platform: Optional[str] = None):
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]
preview_path = f"{basename}.webp"
# Download image file if it is url
if type(image_file_or_url) is str:
image_url = image_file_or_url
# Download file if it is a URL
if type(file_or_url) is str:
url = file_or_url
try:
image_response = requests.get(image_url)
image_response.raise_for_status()
image = Image.open(BytesIO(image_response.content))
image.save(preview_path, "WEBP")
response = requests.get(url)
response.raise_for_status()
# Determine content type from response headers or URL extension
content_type = response.headers.get('content-type', '')
if not content_type:
# Fallback to URL extension detection
content_type = resolve_file_content_type(url) or ''
content = response.content
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}"
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"
image = Image.open(BytesIO(content))
image.save(preview_path, "WEBP")
except Exception as e:
print_error(f"Failed to download image: {e}")
print_error(f"Failed to download preview: {e}")
# Handle uploaded file
else:
# Assert image as file
image_file = image_file_or_url
file_obj = file_or_url
if not isinstance(image_file, web.FileField):
raise RuntimeError("Invalid image file")
if not isinstance(file_obj, web.FileField):
raise RuntimeError("Invalid file")
content_type: str = image_file.content_type
if not content_type.startswith("image/"):
if platform == "huggingface":
# huggingface previewFile content_type='text/plain', not startswith("image/")
return
else:
raise RuntimeError(f"FileTypeError: expected image, got {content_type}")
image = Image.open(image_file.file)
image.save(preview_path, "WEBP")
content_type: str = file_obj.content_type
filename: str = getattr(file_obj, 'filename', '')
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}"
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"
image = Image.open(file_obj.file)
image.save(preview_path, "WEBP")
else:
raise RuntimeError(f"FileTypeError: expected image or video, got {content_type}")
def _get_video_extension_from_url(url: str) -> Optional[str]:
"""Extract video extension from URL."""
from urllib.parse import urlparse
path = urlparse(url).path.lower()
for ext in VIDEO_EXTENSIONS:
if path.endswith(ext):
return ext
return None
def _get_extension_from_content_type(content_type: str) -> Optional[str]:
"""Map content-type to file extension."""
return VIDEO_CONTENT_TYPE_MAP.get(content_type.lower())
def get_model_all_descriptions(model_path: str):
@@ -398,7 +471,7 @@ def rename_model(model_path: str, new_model_path: str):
shutil.move(model_path, new_model_path)
# move preview
previews = get_model_all_images(model_path)
previews = get_model_all_previews(model_path)
for preview in previews:
preview_path = join_path(model_dirname, preview)
preview_name = os.path.splitext(preview)[0]