diff --git a/py/download.py b/py/download.py index 77e7d1b..ebaab42 100644 --- a/py/download.py +++ b/py/download.py @@ -326,7 +326,7 @@ class ModelDownload: try: preview_file = task_data.pop("previewFile", None) - utils.save_model_preview_image(task_path, preview_file, download_platform) + utils.save_model_preview(task_path, preview_file, download_platform) self.set_task_content(task_id, task_data) task_status = TaskStatus( taskId=task_id, diff --git a/py/information.py b/py/information.py index 015ef49..76aade9 100644 --- a/py/information.py +++ b/py/information.py @@ -355,23 +355,17 @@ class Information: @routes.get("/model-manager/preview/{type}/{index}/{filename:.*}") async def read_model_preview(request): """ - Get the file stream of the specified image. + Get the file stream of the specified preview If the file does not exist, no-preview.png is returned. :param type: The type of the model. eg.checkpoints, loras, vae, etc. :param index: The index of the model folders. - :param filename: The filename of the image. + :param filename: The filename of the preview. """ model_type = request.match_info.get("type", None) index = int(request.match_info.get("index", None)) filename = request.match_info.get("filename", None) - content_type = utils.resolve_file_content_type(filename) - - if content_type == "video": - abs_path = utils.get_full_path(model_type, index, filename) - return web.FileResponse(abs_path) - extension_uri = config.extension_uri try: @@ -388,8 +382,16 @@ class Information: if not os.path.isfile(abs_path): abs_path = utils.join_path(extension_uri, "assets", "no-preview.png") - image_data = self.get_image_preview_data(abs_path) - return web.Response(body=image_data.getvalue(), content_type="image/webp") + # Determine content type from the actual file + content_type = utils.resolve_file_content_type(abs_path) + + if content_type == "video": + # Serve video files directly + return web.FileResponse(abs_path) + else: + # Serve image files (WebP or fallback images) + image_data = self.get_image_preview_data(abs_path) + return web.Response(body=image_data.getvalue(), content_type="image/webp") @routes.get("/model-manager/preview/download/{filename}") async def read_download_preview(request): @@ -546,10 +548,10 @@ class Information: model_info = CivitaiModelSearcher().search_by_hash(hash_value) preview_url_list = model_info.get("preview", []) - preview_image_url = preview_url_list[0] if preview_url_list else None - if preview_image_url: - utils.print_debug(f"Save preview image to {abs_image_path}") - utils.save_model_preview_image(abs_model_path, preview_image_url) + preview_url = preview_url_list[0] if preview_url_list else None + if preview_url: + utils.print_debug(f"Save preview to {abs_model_path}") + utils.save_model_preview(abs_model_path, preview_url) description = model_info.get("description", None) if description: diff --git a/py/manager.py b/py/manager.py index 9fe8b6d..a6bfcf7 100644 --- a/py/manager.py +++ b/py/manager.py @@ -134,18 +134,8 @@ class ModelManager: if is_file and extension not in folder_paths.supported_pt_extensions: return None - preview_type = "image" - preview_ext = ".webp" - preview_images = utils.get_model_all_images(entry.path) - if len(preview_images) > 0: - preview_type = "image" - preview_ext = ".webp" - else: - preview_videos = utils.get_model_all_videos(entry.path) - if len(preview_videos) > 0: - preview_type = "video" - preview_ext = f".{preview_videos[0].split('.')[-1]}" - + 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)}" stat = entry.stat() @@ -158,7 +148,6 @@ class ModelManager: "pathIndex": path_index, "sizeBytes": stat.st_size if is_file else 0, "preview": model_preview if is_file else None, - "previewType": preview_type, "createdAt": round(stat.st_ctime_ns / 1000000), "updatedAt": round(stat.st_mtime_ns / 1000000), } @@ -211,10 +200,11 @@ class ModelManager: if "previewFile" in model_data: previewFile = model_data["previewFile"] - if type(previewFile) is str and previewFile == "undefined": - utils.remove_model_preview_image(model_path) - else: - utils.save_model_preview_image(model_path, previewFile) + # Always remove existing preview files first in case the file extension has changed + utils.remove_model_preview(model_path) + # Nothing else to do if the preview file was being removed + if not (type(previewFile) is str and previewFile == "undefined"): + utils.save_model_preview(model_path, previewFile) if "description" in model_data: description = model_data["description"] @@ -236,7 +226,7 @@ class ModelManager: model_dirname = os.path.dirname(model_path) os.remove(model_path) - model_previews = utils.get_model_all_images(model_path) + model_previews = utils.get_model_all_previews(model_path) for preview in model_previews: os.remove(utils.join_path(model_dirname, preview)) diff --git a/py/utils.py b/py/utils.py index c5bdd90..b4f9ceb 100644 --- a/py/utils.py +++ b/py/utils.py @@ -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] diff --git a/src/components/DialogDownload.vue b/src/components/DialogDownload.vue index 0ce440d..f1e48fa 100644 --- a/src/components/DialogDownload.vue +++ b/src/components/DialogDownload.vue @@ -20,7 +20,10 @@ >
- +
+ +
+
@@ -72,11 +75,13 @@ diff --git a/src/components/ResponseFileUpload.vue b/src/components/ResponseFileUpload.vue index 455ff81..e42cd4f 100644 --- a/src/components/ResponseFileUpload.vue +++ b/src/components/ResponseFileUpload.vue @@ -46,7 +46,7 @@ const handleDropFile = (event: DragEvent) => { const handleClick = (event: MouseEvent) => { const input = document.createElement('input') input.type = 'file' - input.accept = 'image/*' + input.accept = 'image/*,video/*' input.onchange = () => { const files = input.files if (files) { diff --git a/src/hooks/explorer.ts b/src/hooks/explorer.ts index 3a496f0..066bb67 100644 --- a/src/hooks/explorer.ts +++ b/src/hooks/explorer.ts @@ -46,7 +46,6 @@ export const useModelExplorer = () => { description: '', metadata: {}, preview: '', - previewType: 'image', type: folder ?? '', isFolder: true, children: [], diff --git a/src/hooks/model.ts b/src/hooks/model.ts index 54b6fba..fa3de56 100644 --- a/src/hooks/model.ts +++ b/src/hooks/model.ts @@ -553,9 +553,11 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => { * Local file url */ const localContent = ref() + const localContentType = ref() const updateLocalContent = async (event: SelectEvent) => { const { files } = event localContent.value = files[0].objectURL + localContentType.value = files[0].type } /** @@ -587,16 +589,13 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => { return content }) - const previewType = computed(() => { - return model.value.previewType - }) - onMounted(() => { registerReset(() => { currentType.value = 'default' defaultContentPage.value = 0 networkContent.value = undefined localContent.value = undefined + localContentType.value = undefined }) registerSubmit((data) => { @@ -606,7 +605,6 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => { const result = { preview, - previewType, typeOptions, currentType, // default value @@ -616,6 +614,7 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => { networkContent, // local file localContent, + localContentType, updateLocalContent, // no preview noPreviewContent, diff --git a/src/types/typings.d.ts b/src/types/typings.d.ts index 66ec0f3..b635158 100644 --- a/src/types/typings.d.ts +++ b/src/types/typings.d.ts @@ -11,7 +11,6 @@ export interface BaseModel { pathIndex: number isFolder: boolean preview: string | string[] - previewType: string description: string metadata: Record } diff --git a/src/utils/media.ts b/src/utils/media.ts new file mode 100644 index 0000000..3f419c7 --- /dev/null +++ b/src/utils/media.ts @@ -0,0 +1,53 @@ +/** + * Media file utility functions + */ + +const VIDEO_EXTENSIONS = [ + '.mp4', + '.webm', + '.mov', + '.avi', + '.mkv', + '.flv', + '.wmv', + '.m4v', + '.ogv', +] + +const VIDEO_HOST_PATTERNS = [ + '/video', // Civitai video URLs often end with /video + 'type=video', // URLs with video type parameter + 'format=video', // URLs with video format parameter + 'video.civitai.com', // Civitai video domain +] + +/** + * Detect if a URL points to a video based on extension or URL patterns + * @param url - The URL to check + * @param localContentType - Optional MIME type for local files + */ +export const isVideoUrl = (url: string, localContentType?: string): boolean => { + if (!url) return false + + // For local files with known MIME type + if (localContentType && localContentType.startsWith('video/')) { + return true + } + + const urlLower = url.toLowerCase() + + // First check if URL ends with a video extension + for (const ext of VIDEO_EXTENSIONS) { + if (urlLower.endsWith(ext)) { + return true + } + } + + // Check if URL contains a video extension anywhere (for complex URLs like Civitai) + if (VIDEO_EXTENSIONS.some((ext) => urlLower.includes(ext))) { + return true + } + + // Check for specific video hosting patterns + return VIDEO_HOST_PATTERNS.some((pattern) => urlLower.includes(pattern)) +}