Feat optimize preview (#156)

* pref: change code structure

* feat(information): support gif preview

* feat(information): support video preview

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()

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),
             }

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
 

src/components/ModelCard.vue

@@ -24,6 +24,21 @@
             ></path>
           </svg>
         </div>
+        <div
+          v-else-if="model.previewType === 'video'"
+          class="h-full w-full p-1 hover:p-0"
+        >
+          <video
+            class="h-full w-full object-cover"
+            playsinline
+            autoplay
+            loop
+            disablepictureinpicture
+            preload="none"
+          >
+            <source :src="preview" />
+          </video>
+        </div>
         <div v-else class="h-full w-full p-1 hover:p-0">
           <img class="h-full w-full rounded-lg object-cover" :src="preview" />
         </div>

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` })"
       >
-        <ResponseImage :src="preview" :error="noPreviewContent"></ResponseImage>
+        <div v-if="previewType === 'video'" class="h-full w-full p-1 hover:p-0">
+          <video
+            class="h-full w-full object-cover"
+            playsinline
+            autoplay
+            loop
+            disablepictureinpicture
+            preload="none"
+          >
+            <source :src="preview" />
+          </video>
+        </div>
+
+        <ResponseImage
+          v-else
+          :src="preview"
+          :error="noPreviewContent"
+        ></ResponseImage>
 
         <Carousel
           v-if="defaultContent.length > 1"
@@ -95,6 +112,7 @@ const { cardWidth } = useConfig()
 
 const {
   preview,
+  previewType,
   typeOptions,
   currentType,
   defaultContent,

src/hooks/model.ts

@@ -579,6 +579,10 @@ export const useModelPreviewEditor = (formInstance: ModelFormInstance) => {
     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

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<string, string>
 }