Preview image improvements.

- Model Tab grid receives smaller previews from server.
- Attempted to make PIL image `info` serializable for previews.
- Get full size previews from Civitai.
  - Note, the Civitai server may return nothing for the image id. (External bug?)
- Support downloading previews from https://civitai.com/images/
- Lazy Loading in Model Tab.

__init__.py

@@ -198,6 +198,26 @@ def server_rules():
 server_settings = config_loader.yaml_load(server_settings_uri, server_rules())
 config_loader.yaml_save(server_settings_uri, server_rules(), server_settings)
 
+
+def get_def_headers(url=""):
+    def_headers = {
+        "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148",
+    }
+
+    if url.startswith("https://civitai.com/"):
+        api_key = server_settings["civitai_api_key"]
+        if (api_key != ""):
+            def_headers["Authorization"] = f"Bearer {api_key}"
+            url += "&" if "?" in url else "?" # not the most robust solution
+            url += f"token={api_key}" # TODO: Authorization didn't work in the header
+    elif url.startswith("https://huggingface.co/"):
+        api_key = server_settings["huggingface_api_key"]
+        if api_key != "":
+            def_headers["Authorization"] = f"Bearer {api_key}"
+    
+    return def_headers
+
+
 @server.PromptServer.instance.routes.get("/model-manager/settings/load")
 async def load_ui_settings(request):
     rules = ui_rules()
@@ -218,34 +238,105 @@ async def save_ui_settings(request):
     })
 
 
+from PIL import Image, TiffImagePlugin
+from PIL.PngImagePlugin import PngInfo
+def PIL_cast_serializable(v):
+    # source: https://github.com/python-pillow/Pillow/issues/6199#issuecomment-1214854558
+    if isinstance(v, TiffImagePlugin.IFDRational):
+        return float(v)
+    elif isinstance(v, tuple):
+        return tuple(PIL_cast_serializable(t) for t in v)
+    elif isinstance(v, bytes):
+        return v.decode(errors="replace")
+    elif isinstance(v, dict):
+        for kk, vv in v.items():
+            v[kk] = PIL_cast_serializable(vv)
+        return v
+    else:
+        return v
+
+
+def get_safetensors_image_bytes(path):
+    if not os.path.isfile(path):
+        raise RuntimeError("Path was invalid!")
+    header = get_safetensor_header(path)
+    metadata = header.get("__metadata__", None)
+    if metadata is None:
+        return None
+    thumbnail = metadata.get("modelspec.thumbnail", None)
+    if thumbnail is None:
+        return None
+    image_data = thumbnail.split(',')[1]
+    return base64.b64decode(image_data)
+
+
 @server.PromptServer.instance.routes.get("/model-manager/preview/get")
 async def get_model_preview(request):
     uri = request.query.get("uri")
-
     image_path = no_preview_image
     image_type = "png"
-    image_data = None
     if uri != "no-preview":
         sep = os.path.sep
         uri = uri.replace("/" if sep == "\\" else "/", sep)
         path, _ = search_path_to_system_path(uri)
         head, extension = split_valid_ext(path, preview_extensions)
         if os.path.exists(path):
-            image_type = extension.rsplit(".", 1)[1]
             image_path = path
-        elif os.path.exists(head) and head.endswith(".safetensors"):
             image_type = extension.rsplit(".", 1)[1]
-            header = get_safetensor_header(head)
+        elif os.path.exists(head) and head.endswith(".safetensors"):
-            metadata = header.get("__metadata__", None)
+            image_path = head
-            if metadata is not None:
+            image_type = extension.rsplit(".", 1)[1]
-                thumbnail = metadata.get("modelspec.thumbnail", None)
-                if thumbnail is not None:
-                    image_data = thumbnail.split(',')[1]
-                    image_data = base64.b64decode(image_data)
 
-    if image_data == None:
+    w = request.query.get("width")
-        with open(image_path, "rb") as file:
+    h = request.query.get("height")
-            image_data = file.read()
+    try:
+        w = int(w)
+        if w < 1:
+            w = None
+    except:
+        w = None
+    try:
+        h = int(h)
+        if w < 1:
+            h = None
+    except:
+        h = None
+
+    image_data = None
+    if w is None and h is None: # full size
+        if image_path.endswith(".safetensors"):
+            image_data = get_safetensors_image_bytes(image_path)
+        else:
+            with open(image_path, "rb") as image:
+                image_data = image.read()
+    else:
+        if image_path.endswith(".safetensors"):
+            image_data = get_safetensors_image_bytes(image_path)
+            fp = io.BytesIO(image_data)
+        else:
+            fp = image_path
+
+        with Image.open(fp) as image:
+            w0, h0 = image.size
+            if w is None:
+                w = (h * w0) // h0
+            elif h is None:
+                h = (w * h0) // w0
+
+            exif = image.getexif()
+
+            metadata = None
+            if len(image.info) > 0:
+                metadata = PngInfo()
+                for (key, value) in image.info.items():
+                    value_str = str(PIL_cast_serializable(value)) # not sure if this is correct (sometimes includes exif)
+                    metadata.add_text(key, value_str)
+
+            image.thumbnail((w, h))
+
+            image_bytes = io.BytesIO()
+            image.save(image_bytes, format=image.format, exif=exif, pnginfo=metadata)
+            image_data = image_bytes.getvalue()
 
     return web.Response(body=image_data, content_type="image/" + image_type)
 
@@ -268,7 +359,30 @@ def download_model_preview(formdata):
 
     image = formdata.get("image", None)
     if type(image) is str:
-        _, image_extension = split_valid_ext(image, image_extensions) # TODO: doesn't work for https://civitai.com/images/...
+        civitai_image_url = "https://civitai.com/images/"
+        if image.startswith(civitai_image_url):
+            image_id = re.search(r"^\d+", image[len(civitai_image_url):]).group(0)
+            image_id = str(int(image_id))
+            image_info_url = f"https://civitai.com/api/v1/images?imageId={image_id}"
+            def_headers = get_def_headers(image_info_url)
+            response = requests.get(
+                url=image_info_url, 
+                stream=False, 
+                verify=False, 
+                headers=def_headers, 
+                proxies=None, 
+                allow_redirects=False, 
+            )
+            if response.ok:
+                content_type = response.headers.get("Content-Type")
+                info = response.json()
+                items = info["items"]
+                if len(items) == 0:
+                    raise RuntimeError("Civitai /api/v1/images returned 0 items!")
+                image = items[0]["url"]
+            else:
+                raise RuntimeError("Bad response from api/v1/images!")
+        _, image_extension = split_valid_ext(image, image_extensions)
         if image_extension == "":
             raise ValueError("Invalid image type!")
         image_path = path_without_extension + image_extension
@@ -474,21 +588,15 @@ def download_file(url, filename, overwrite):
 
     filename_temp = filename + ".download"
 
-    def_headers = {
+    def_headers = get_def_headers(url)
-        "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148",
+    rh = requests.get(
-    }
+        url=url, 
-
+        stream=True, 
-    if url.startswith("https://civitai.com/"):
+        verify=False, 
-        api_key = server_settings["civitai_api_key"]
+        headers=def_headers, 
-        if (api_key != ""):
+        proxies=None, 
-            def_headers["Authorization"] = f"Bearer {api_key}"
+        allow_redirects=False, 
-            url += "&" if "?" in url else "?" # not the most robust solution
+    )
-            url += f"token={api_key}" # TODO: Authorization didn't work in the header
-    elif url.startswith("https://huggingface.co/"):
-        api_key = server_settings["huggingface_api_key"]
-        if api_key != "":
-            def_headers["Authorization"] = f"Bearer {api_key}"
-    rh = requests.get(url=url, stream=True, verify=False, headers=def_headers, proxies=None, allow_redirects=False)
     if not rh.ok:
         raise ValueError(
             "Unable to download! Request header status code: " + 
@@ -501,8 +609,16 @@ def download_file(url, filename, overwrite):
 
     headers = {"Range": "bytes=%d-" % downloaded_size}
     headers["User-Agent"] = def_headers["User-Agent"]
-    
+    headers["Authorization"] = def_headers.get("Authorization", None)
-    r = requests.get(url=url, stream=True, verify=False, headers=headers, proxies=None, allow_redirects=False)
+
+    r = requests.get(
+        url=url, 
+        stream=True, 
+        verify=False, 
+        headers=headers, 
+        proxies=None, 
+        allow_redirects=False, 
+    )
     if rh.status_code == 307 and r.status_code == 307:
         # Civitai redirect
         redirect_url = r.content.decode("utf-8")

web/model-manager.js

@@ -112,18 +112,28 @@ class SearchPath {
 /**
  * @param {string | undefined} [searchPath=undefined]
  * @param {string | undefined} [dateImageModified=undefined]
- *
+ * @param {string | undefined} [width=undefined]
+ * @param {string | undefined} [height=undefined]
  * @returns {string}
  */
-function imageUri(imageSearchPath = undefined, dateImageModified = undefined) {
+function imageUri(imageSearchPath = undefined, dateImageModified = undefined, width = undefined, height = undefined) {
     const path = imageSearchPath ?? "no-preview";
     const date = dateImageModified;
     let uri = `/model-manager/preview/get?uri=${path}`;
+    if (width !== undefined && width !== null) {
+        uri += `&width=${width}`;
+    }
+    if (height !== undefined && height !== null) {
+        uri += `&height=${height}`;
+    }
     if (date !== undefined && date !== null) {
         uri += `&v=${date}`;
     }
     return uri;
 }
+const PREVIEW_NONE_URI = imageUri();
+const PREVIEW_THUMBNAIL_WIDTH = 320;
+const PREVIEW_THUMBNAIL_HEIGHT = 480;
 
 /**
  * @param {(...args) => void} callback
@@ -334,31 +344,54 @@ class ImageSelect {
     /** @type {string} */
     #name = null;
     
-    /** @returns {string|File} */
+    /** @returns {Promise<string> | Promise<File>} */
-    getImage() {
+    async getImage() {
         const name = this.#name;
         const value = document.querySelector(`input[name="${name}"]:checked`).value;
         const elements = this.elements;
         switch (value) {
             case this.#PREVIEW_DEFAULT:
                 const children = elements.defaultPreviews.children;
-                const noImage = imageUri();
+                const noImage = PREVIEW_NONE_URI;
+                let url = "";
                 for (let i = 0; i < children.length; i++) {
                     const child = children[i];
                     if (child.style.display !== "none" && 
                         child.nodeName === "IMG" && 
                         !child.src.endsWith(noImage)
                     ) {
-                        return child.src;
+                        url = child.src;
                     }
                 }
-                return "";
+                if (url.startsWith(Civitai.imageUrlPrefix())) {
+                    url = await Civitai.getFullSizeImageUrl(url).catch((err) => {
+                        console.warn(err);
+                        return url;
+                    });
+                }
+                return url;
             case this.#PREVIEW_URL:
-                return elements.customUrl.value;
+                const value = elements.customUrl.value;
+                if (value.startsWith(Civitai.imagePostUrlPrefix())) {
+                    try {
+                        const imageInfo = await Civitai.getImageInfo(value);
+                        const items = imageInfo["items"];
+                        if (items.length === 0) {
+                            console.warn("Civitai /api/v1/images returned 0 items.");
+                            return value;
+                        }
+                        return items[0]["url"];
+                    }
+                    catch (error) {
+                        console.error("Failed to get image info from Civitai!", error);
+                        return value;
+                    }
+                }
+                return value;
             case this.#PREVIEW_UPLOAD:
                 return elements.uploadFile.files[0] ?? "";
             case this.#PREVIEW_NONE:
-                return imageUri();
+                return PREVIEW_NONE_URI;
         }
         return "";
     }
@@ -382,7 +415,7 @@ class ImageSelect {
                 }
             }
             else {
-                el.src = imageUri();
+                el.src = PREVIEW_NONE_URI;
             }
         });
         this.checkDefault();
@@ -448,19 +481,19 @@ class ImageSelect {
      */
     constructor(radioGroupName, defaultPreviews = []) {
         if (defaultPreviews === undefined | defaultPreviews === null | defaultPreviews.length === 0) {
-            defaultPreviews = [imageUri()];
+            defaultPreviews = [PREVIEW_NONE_URI];
         }
         this.#name = radioGroupName;
         
         const el_defaultUri = $el("div", {
             $: (el) => (this.elements.defaultUrl = el),
             style: { display: "none" },
-            "data-noimage": imageUri(),
+            "data-noimage": PREVIEW_NONE_URI,
         });
         
         const el_defaultPreviewNoImage = $el("img", {
             $: (el) => (this.elements.defaultPreviewNoImage = el),
-            src: imageUri(),
+            src: PREVIEW_NONE_URI,
             style: { display: "none" },
             loading: "lazy",
         });
@@ -478,7 +511,7 @@ class ImageSelect {
                     style: { display: "none" },
                     loading: "lazy",
                     onerror: (e) => {
-                        e.target.src = el_defaultUri.dataset.noimage ?? imageUri();
+                        e.target.src = el_defaultUri.dataset.noimage ?? PREVIEW_NONE_URI;
                     },
                 });
             });
@@ -490,10 +523,10 @@ class ImageSelect {
         
         const el_uploadPreview = $el("img", {
             $: (el) => (this.elements.uploadPreview = el),
-            src: imageUri(),
+            src: PREVIEW_NONE_URI,
             style: { display : "none" },
             onerror: (e) => {
-                e.target.src = el_defaultUri.dataset.noimage ?? imageUri();
+                e.target.src = el_defaultUri.dataset.noimage ?? PREVIEW_NONE_URI;
             },
         });
         const el_uploadFile = $el("input", {
@@ -520,10 +553,10 @@ class ImageSelect {
         
         const el_customUrlPreview = $el("img", {
             $: (el) => (this.elements.customUrlPreview = el),
-            src: imageUri(),
+            src: PREVIEW_NONE_URI,
             style: { display: "none" },
             onerror: (e) => {
-                e.target.src = el_defaultUri.dataset.noimage ?? imageUri();
+                e.target.src = el_defaultUri.dataset.noimage ?? PREVIEW_NONE_URI;
             },
         });
         const el_customUrl = $el("input.search-text-area", {
@@ -540,8 +573,28 @@ class ImageSelect {
             el_customUrl,
             $el("button.icon-button", {
                 textContent: "🔍︎",
-                onclick: (e) => {
+                onclick: async (e) => {
-                    el_customUrlPreview.src = el_customUrl.value;
+                    const value = el_customUrl.value;
+                    if (value.startsWith(Civitai.imagePostUrlPrefix())) {
+                        el_customUrlPreview.src = await Civitai.getImageInfo(value)
+                        .then((imageInfo) => {
+                            const items = imageInfo["items"];
+                            if (items.length > 0) {
+                                return items[0]["url"];
+                            }
+                            else {
+                                console.warn("Civitai /api/v1/images returned 0 items.");
+                                return value;
+                            }
+                        })
+                        .catch((error) => {
+                            console.error("Failed to get image info from Civitai!", error);
+                            return value;
+                        });
+                    }
+                    else {
+                        el_customUrlPreview.src = value;
+                    }
                 },
             }),
         ]);
@@ -1534,8 +1587,14 @@ class ModelGrid {
                 );
                 return $el("div.item", {}, [
                     $el("img.model-preview", {
-                        src: imageUri(previewInfo?.path, previewInfo?.dateModified),
+                        src: imageUri(
+                            previewInfo?.path, 
+                            previewInfo?.dateModified, 
+                            PREVIEW_THUMBNAIL_WIDTH, 
+                            PREVIEW_THUMBNAIL_HEIGHT, 
+                        ),
                         draggable: false,
+                        loading: "lazy",
                     }),
                     $el("div.model-preview-overlay", {
                         ondragend: (e) => dragAdd(e),
@@ -1674,8 +1733,8 @@ class ModelInfoView {
                     e.target.disabled = true;
                     const container = this.elements.info;
                     const path = container.dataset.path;
-                    const imageUrl = previewSelect.getImage();
+                    const imageUrl = await previewSelect.getImage();
-                    if (imageUrl === imageUri()) {
+                    if (imageUrl === PREVIEW_NONE_URI) {
                         const encodedPath = encodeURIComponent(path);
                         updatedPreview = await request(
                             `/model-manager/preview/delete?path=${encodedPath}`,
@@ -1713,7 +1772,7 @@ class ModelInfoView {
                     if (updatedPreview) {
                         updateModels();
                         const previewSelect = this.previewSelect;
-                        previewSelect.elements.defaultUrl.dataset.noimage = imageUri();
+                        previewSelect.elements.defaultUrl.dataset.noimage = PREVIEW_NONE_URI;
                         previewSelect.resetModelInfoPreview();
                         this.element.style.display = "none";
                     }
@@ -1936,7 +1995,7 @@ class ModelInfoView {
             defaultUrl.dataset.noimage = imageUri(imagePath, imageDateModified);
         }
         else {
-            defaultUrl.dataset.noimage = imageUri();
+            defaultUrl.dataset.noimage = PREVIEW_NONE_URI;
         }
         previewSelect.resetModelInfoPreview();
         const setPreviewButton = this.elements.setPreviewButton;
@@ -2111,8 +2170,6 @@ class Civitai {
     }
     
     /**
-     * 
-     *
      * @param {string} stringUrl - Model url.
      *
      * @returns {Promise<Object>} - Download information for the given url.
@@ -2182,6 +2239,73 @@ class Civitai {
             return {};
         }
     }
+    
+    /**
+     * @returns {string}
+     */
+    static imagePostUrlPrefix() {
+        return "https://civitai.com/images/";
+    }
+    
+    /**
+     * @returns {string}
+     */
+    static imageUrlPrefix() {
+        return "https://image.civitai.com/";
+    }
+    
+    /**
+     * @param {string} stringUrl - https://civitai.com/images/{imageId}.
+     *
+     * @returns {Promise<Object>} - Image information.
+     */
+    static async getImageInfo(stringUrl) {
+        const imagePostUrlPrefix = Civitai.imagePostUrlPrefix();
+        if (!stringUrl.startsWith(imagePostUrlPrefix)) {
+            return {};
+        }
+        const id = stringUrl.substring(imagePostUrlPrefix.length).match(/^\d+/)[0];
+        const url = `https://civitai.com/api/v1/images?imageId=${id}`;
+        try {
+            return await request(url);
+        }
+        catch (error) {
+            console.error("Failed to get image info from Civitai!", error);
+            return {};
+        }
+    }
+    
+    /**
+     * @param {string} stringUrl - https://image.civitai.com/...
+     *
+     * @returns {Promise<string>}
+     */
+    static async getFullSizeImageUrl(stringUrl) {
+        const imageUrlPrefix = Civitai.imageUrlPrefix();
+        if (!stringUrl.startsWith(imageUrlPrefix)) {
+            return "";
+        }
+        const i0 = stringUrl.lastIndexOf("/");
+        const i1 = stringUrl.lastIndexOf(".");
+        if (i0 === -1 || i1 === -1) {
+            return "";
+        }
+        const id = parseInt(stringUrl.substring(i0 + 1, i1)).toString();
+        const url = `https://civitai.com/api/v1/images?imageId=${id}`;
+        try {
+            const imageInfo = await request(url);
+            const items = imageInfo["items"];
+            if (items.length === 0) {
+                console.warn("Civitai /api/v1/images returned 0 items.");
+                return stringUrl;
+            }
+            return items[0]["url"];
+        }
+        catch (error) {
+            console.error("Failed to get image info from Civitai!", error);
+            return stringUrl;
+        }
+    }
 }
 
 class HuggingFace {
@@ -2427,8 +2551,8 @@ class DownloadTab {
                                     }) ?? "";
                                     return name + ext;
                                 })());
-                                const image = downloadPreviewSelect.getImage();
+                                const image = await downloadPreviewSelect.getImage();
-                                formData.append("image", image === imageUri() ? "" : image);
+                                formData.append("image", image === PREVIEW_NONE_URI ? "" : image);
                                 formData.append("overwrite", this.elements.overwrite.checked);
                                 e.target.disabled = true;
                                 const [success, resultText] = await request(