From f015767085d40c128ad683eea297d75c3cf30978 Mon Sep 17 00:00:00 2001 From: Christian Bastian <80225746+cdb-boop@users.noreply.github.com> Date: Mon, 1 Apr 2024 00:12:09 -0400 Subject: [PATCH] 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 | 180 ++++++++++++++++++++++++++++++++++-------- web/model-manager.js | 182 ++++++++++++++++++++++++++++++++++++------- 2 files changed, 301 insertions(+), 61 deletions(-) diff --git a/__init__.py b/__init__.py index a664257..e6301a2 100644 --- a/__init__.py +++ b/__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) - metadata = header.get("__metadata__", None) - if metadata is not None: - thumbnail = metadata.get("modelspec.thumbnail", None) - if thumbnail is not None: - image_data = thumbnail.split(',')[1] - image_data = base64.b64decode(image_data) + elif os.path.exists(head) and head.endswith(".safetensors"): + image_path = head + image_type = extension.rsplit(".", 1)[1] - if image_data == None: - with open(image_path, "rb") as file: - image_data = file.read() + w = request.query.get("width") + h = request.query.get("height") + 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 = { - "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}" - rh = requests.get(url=url, stream=True, verify=False, headers=def_headers, proxies=None, allow_redirects=False) + def_headers = get_def_headers(url) + 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"] - - r = requests.get(url=url, stream=True, verify=False, headers=headers, proxies=None, allow_redirects=False) + headers["Authorization"] = def_headers.get("Authorization", None) + + 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") diff --git a/web/model-manager.js b/web/model-manager.js index 138cd8f..595d0c1 100644 --- a/web/model-manager.js +++ b/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} */ - getImage() { + /** @returns {Promise | Promise} */ + 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) => { - el_customUrlPreview.src = el_customUrl.value; + onclick: async (e) => { + 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(); - if (imageUrl === imageUri()) { + const imageUrl = await previewSelect.getImage(); + 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} - 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} - 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} + */ + 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(); - formData.append("image", image === imageUri() ? "" : image); + const image = await downloadPreviewSelect.getImage(); + formData.append("image", image === PREVIEW_NONE_URI ? "" : image); formData.append("overwrite", this.elements.overwrite.checked); e.target.disabled = true; const [success, resultText] = await request(