Add local preview file functionality

__init__.py

@@ -211,6 +211,10 @@ async def api_view_file(request):
     filename = request.query.get("filename", None)
     if not filename:
         return web.Response(status=404)
+    
+    # validation for security: prevent accessing arbitrary path
+    if filename[0] == '/' or '..' in filename:
+        return web.Response(status=400)
 
     output_path = folder_paths.get_output_directory()
     file_path = os.path.join(output_path, filename)

service/model_manager/model_list.py

@@ -9,6 +9,7 @@ import threading
 import os
 import urllib.request
 import json
+from .model_preview import preview_file 
 
 # The path to the file where file_hash_dict will be saved
 FILE_HASH_DICT_FOLDER_PATH = os.path.join(os.path.dirname(__file__),"../../hash")
@@ -46,7 +47,7 @@ def process_file(folder, file):
     else:
         file_hash = None  # placeholder for hash
     [model_name, model_extension] = os.path.splitext(file)
-    return {"model_name": model_name, "model_type": folder, "model_extension": model_extension, "file_hash": file_hash}
+    return {"model_name": model_name, "model_type": folder, "model_extension": model_extension, "file_hash": file_hash, "preview": preview_file(file_path)}
 
 def populate_file_hash_dict():
     global populate_done, file_list

service/model_manager/model_preview.py

@@ -0,0 +1,30 @@
+import os
+from PIL import Image
+import base64
+from io import BytesIO
+
+def preview_file(filename: str):
+    preview_exts = [".jpg", ".png", ".jpeg", ".gif"]
+    preview_exts = [*preview_exts, *[".preview" + x for x in preview_exts]]
+    for ext in preview_exts:
+        path = os.path.splitext(filename)[0] + ext
+        if os.path.exists(path):
+            # because ComfyUI has extra model path feature
+            # the path might not be relative to the ComfyUI root
+            # so instead of returning the path, we return the image data directly, to avoid security issues
+            with Image.open(path) as img:
+                # If the image is too large, resize it
+                if img.width > 128 and img.height > 178:
+                    # Calculate new width to maintain aspect ratio
+                    width = int(img.width * 178 / img.height)
+                    # Resize the image
+                    img = img.resize((width, 178))
+                img = img.convert("RGB")
+                # Save the image to a BytesIO object
+                buffer = BytesIO()
+                img.save(buffer, format="JPEG", quality=85)
+                # Get the base64 string
+                img_base64 = base64.b64encode(buffer.getvalue()).decode()
+                # Return the base64 string
+                return f"data:image/jpeg;base64, {img_base64}"
+    return None

ui/src/model-manager/models-list-drawer/ModelItem.tsx

@@ -53,9 +53,10 @@ export function ModelItem({ data }: Props) {
     const model = await indexdb.models.get(
       data.model_name + "@" + data.model_type,
     );
+    if (data.preview) setUrl(data.preview);
     if (model != null) {
       setModel(model);
-      model.imageUrl?.length && setUrl(model.imageUrl);
+      !data.preview && model.imageUrl?.length && setUrl(model.imageUrl);
     } else if (data.file_hash != null) {
       try {
         const url = `https://civitai.com/api/v1/model-versions/by-hash/${data.file_hash}`;
@@ -71,7 +72,7 @@ export function ModelItem({ data }: Props) {
           const sfwImage = json.images.find((i) => i.nsfw === "None");
           image_url = sfwImage?.url;
         }
-        image_url && setUrl(image_url);
+        !data.preview && image_url && setUrl(image_url);
 
         indexdb.models.add({
           id: data.model_name + "@" + data.model_type,

ui/src/model-manager/types.ts

@@ -132,4 +132,5 @@ export interface ModelsListRespItem {
   model_extension: string;
   model_type: string;
   file_hash?: string;
+  preview?: string;
 }