diff --git a/__init__.py b/__init__.py index a1ffb74..7696d39 100644 --- a/__init__.py +++ b/__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) diff --git a/service/model_manager/model_list.py b/service/model_manager/model_list.py index 822461d..9b224f4 100644 --- a/service/model_manager/model_list.py +++ b/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 diff --git a/service/model_manager/model_preview.py b/service/model_manager/model_preview.py new file mode 100644 index 0000000..cbc5771 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ui/src/model-manager/models-list-drawer/ModelItem.tsx b/ui/src/model-manager/models-list-drawer/ModelItem.tsx index 9ebb1f1..7792ccb 100644 --- a/ui/src/model-manager/models-list-drawer/ModelItem.tsx +++ b/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, diff --git a/ui/src/model-manager/types.ts b/ui/src/model-manager/types.ts index d4d35bd..48754a6 100644 --- a/ui/src/model-manager/types.ts +++ b/ui/src/model-manager/types.ts @@ -132,4 +132,5 @@ export interface ModelsListRespItem { model_extension: string; model_type: string; file_hash?: string; + preview?: string; }