From 7378a7deae270b156c3be3b5e8321b1324655f36 Mon Sep 17 00:00:00 2001 From: Hayden <48267247+hayden-fr@users.noreply.github.com> Date: Mon, 3 Mar 2025 14:50:06 +0800 Subject: [PATCH] Feat optimize preview (#156) * pref: change code structure * feat(information): support gif preview * feat(information): support video preview --- py/information.py | 110 ++++++++++++++++++++++++++------ py/manager.py | 15 ++++- py/utils.py | 31 +++++++++ src/components/ModelCard.vue | 15 +++++ src/components/ModelPreview.vue | 20 +++++- src/hooks/model.ts | 5 ++ src/types/typings.d.ts | 1 + 7 files changed, 174 insertions(+), 23 deletions(-) diff --git a/py/information.py b/py/information.py index 16aa0ea..6593e7e 100644 --- a/py/information.py +++ b/py/information.py @@ -1,14 +1,22 @@ import os import re +import math import yaml import requests import markdownify +import folder_paths + +from aiohttp import web from abc import ABC, abstractmethod from urllib.parse import urlparse, parse_qs +from PIL import Image +from io import BytesIO + from . import utils +from . import config class ModelSearcher(ABC): @@ -282,25 +290,6 @@ class HuggingfaceModelSearcher(ModelSearcher): return _filter_tree_files -def get_model_searcher_by_url(url: str) -> ModelSearcher: - parsed_url = urlparse(url) - host_name = parsed_url.hostname - if host_name == "civitai.com": - return CivitaiModelSearcher() - elif host_name == "huggingface.co": - return HuggingfaceModelSearcher() - return UnknownWebsiteSearcher() - - -import folder_paths - - -from . import config - - -from aiohttp import web - - class Information: def add_routes(self, routes): @@ -347,18 +336,30 @@ class Information: 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: folders = folder_paths.get_folder_paths(model_type) base_path = folders[index] abs_path = utils.join_path(base_path, filename) + preview_name = utils.get_model_preview_name(abs_path) + if preview_name: + dir_name = os.path.dirname(abs_path) + abs_path = utils.join_path(dir_name, preview_name) except: abs_path = extension_uri if not os.path.isfile(abs_path): abs_path = utils.join_path(extension_uri, "assets", "no-preview.png") - return web.FileResponse(abs_path) + + 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): @@ -373,11 +374,69 @@ class Information: return web.FileResponse(preview_path) + def get_image_preview_data(self, filename: str): + with Image.open(filename) as img: + max_size = 1024 + original_format = img.format + + exif_data = img.info.get("exif") + icc_profile = img.info.get("icc_profile") + + if getattr(img, "is_animated", False) and img.n_frames > 1: + total_frames = img.n_frames + step = max(1, math.ceil(total_frames / 30)) + + frames, durations = [], [] + + for frame_idx in range(0, total_frames, step): + img.seek(frame_idx) + frame = img.copy() + frame.thumbnail((max_size, max_size), Image.Resampling.NEAREST) + + frames.append(frame) + durations.append(img.info.get("duration", 100) * step) + + save_args = { + "format": "WEBP", + "save_all": True, + "append_images": frames[1:], + "duration": durations, + "loop": 0, + "quality": 80, + "method": 0, + "allow_mixed": False, + } + + if exif_data: + save_args["exif"] = exif_data + + if icc_profile: + save_args["icc_profile"] = icc_profile + + img_byte_arr = BytesIO() + frames[0].save(img_byte_arr, **save_args) + img_byte_arr.seek(0) + return img_byte_arr + + img.thumbnail((max_size, max_size), Image.Resampling.BICUBIC) + + img_byte_arr = BytesIO() + save_args = {"format": "WEBP", "quality": 80} + + if exif_data: + save_args["exif"] = exif_data + if icc_profile: + save_args["icc_profile"] = icc_profile + + img.save(img_byte_arr, **save_args) + img_byte_arr.seek(0) + return img_byte_arr + def fetch_model_info(self, model_page: str): if not model_page: return [] - model_searcher = get_model_searcher_by_url(model_page) + model_searcher = self.get_model_searcher_by_url(model_page) result = model_searcher.search_by_url(model_page) return result @@ -435,3 +494,12 @@ class Information: utils.print_error(f"Failed to download model info for {abs_model_path}: {e}") utils.print_debug("Completed scan model information.") + + def get_model_searcher_by_url(self, url: str) -> ModelSearcher: + parsed_url = urlparse(url) + host_name = parsed_url.hostname + if host_name == "civitai.com": + return CivitaiModelSearcher() + elif host_name == "huggingface.co": + return HuggingfaceModelSearcher() + return UnknownWebsiteSearcher() diff --git a/py/manager.py b/py/manager.py index f7ddc15..a4642cb 100644 --- a/py/manager.py +++ b/py/manager.py @@ -134,7 +134,19 @@ class ModelManager: if is_file and extension not in folder_paths.supported_pt_extensions: return None - model_preview = f"/model-manager/preview/{folder}/{path_index}/{relative_path.replace(extension, '.webp')}" + 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]}" + + model_preview = f"/model-manager/preview/{folder}/{path_index}/{relative_path.replace(extension, preview_ext)}" stat = entry.stat() return { @@ -146,6 +158,7 @@ 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), } diff --git a/py/utils.py b/py/utils.py index eed5551..5114b17 100644 --- a/py/utils.py +++ b/py/utils.py @@ -8,6 +8,7 @@ import requests import traceback import configparser import functools +import mimetypes import comfy.utils import folder_paths @@ -149,6 +150,20 @@ def resolve_model_base_paths() -> dict[str, list[str]]: return model_base_paths +def resolve_file_content_type(filename: str): + extension_mimetypes_cache = folder_paths.extension_mimetypes_cache + extension = filename.split(".")[-1] + if extension not in extension_mimetypes_cache: + mime_type, _ = mimetypes.guess_type(filename, strict=False) + if not mime_type: + return None + content_type = mime_type.split("/")[0] + extension_mimetypes_cache[extension] = content_type + else: + content_type = extension_mimetypes_cache[extension] + return content_type + + def get_full_path(model_type: str, path_index: int, filename: str): """ Get the absolute path in the model type through string concatenation. @@ -266,6 +281,22 @@ def get_model_preview_name(model_path: str): return images[0] if len(images) > 0 else "no-preview.png" +def get_model_all_videos(model_path: str): + 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 + + from PIL import Image from io import BytesIO diff --git a/src/components/ModelCard.vue b/src/components/ModelCard.vue index 205f4b4..b3af905 100644 --- a/src/components/ModelCard.vue +++ b/src/components/ModelCard.vue @@ -24,6 +24,21 @@ > +
+ +
diff --git a/src/components/ModelPreview.vue b/src/components/ModelPreview.vue index 215b503..b3a3b07 100644 --- a/src/components/ModelPreview.vue +++ b/src/components/ModelPreview.vue @@ -5,7 +5,24 @@ class="relative mx-auto w-full overflow-hidden rounded-lg preview-aspect" :style="$sm({ width: `${cardWidth}px` })" > - +
+ +
+ + { return content }) + const previewType = computed(() => { + return model.value.previewType + }) + onMounted(() => { registerReset(() => { currentType.value = 'default' @@ -594,6 +598,7 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => { const result = { preview, + previewType, typeOptions, currentType, // default value diff --git a/src/types/typings.d.ts b/src/types/typings.d.ts index 35162ca..6a554cd 100644 --- a/src/types/typings.d.ts +++ b/src/types/typings.d.ts @@ -11,6 +11,7 @@ export interface BaseModel { pathIndex: number isFolder: boolean preview: string | string[] + previewType: string description: string metadata: Record }