TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Preview image improvements.

Browse Source 
- 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.
This commit is contained in:
Christian Bastian
2024-04-01 00:12:09 -04:00
parent a15ec0006e
commit f015767085
2 changed files with 301 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
__init__.py
View File

@@ -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")

182
web/model-manager.js
View File

@@ -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<string> | Promise<File>} */
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<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();
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(