diff --git a/extensions-builtin/Lora/network.py b/extensions-builtin/Lora/network.py index 412864292..20f8df3d4 100644 --- a/extensions-builtin/Lora/network.py +++ b/extensions-builtin/Lora/network.py @@ -29,7 +29,6 @@ class NetworkOnDisk: def read_metadata(): metadata = sd_models.read_metadata_from_safetensors(filename) - metadata.pop('ssmd_cover_images', None) # those are cover images, and they are too big to display in UI as text return metadata diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py index 66d15dd05..b627f7dc2 100644 --- a/extensions-builtin/Lora/ui_extra_networks_lora.py +++ b/extensions-builtin/Lora/ui_extra_networks_lora.py @@ -31,7 +31,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): "name": name, "filename": lora_on_disk.filename, "shorthash": lora_on_disk.shorthash, - "preview": self.find_preview(path), + "preview": self.find_preview(path) or self.find_embedded_preview(path, name, lora_on_disk.metadata), "description": self.find_description(path), "search_terms": search_terms, "local_preview": f"{path}.{shared.opts.samples_format}", diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index f4627ce8d..087cb8a05 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -1,6 +1,8 @@ import functools import os.path import urllib.parse +from base64 import b64decode +from io import BytesIO from pathlib import Path from typing import Optional, Union from dataclasses import dataclass @@ -11,6 +13,7 @@ import gradio as gr import json import html from fastapi.exceptions import HTTPException +from PIL import Image from modules.infotext_utils import image_from_url_text @@ -108,6 +111,31 @@ def fetch_file(filename: str = ""): return FileResponse(filename, headers={"Accept-Ranges": "bytes"}) +def fetch_cover_images(page: str = "", item: str = "", index: int = 0): + from starlette.responses import Response + + page = next(iter([x for x in extra_pages if x.name == page]), None) + if page is None: + raise HTTPException(status_code=404, detail="File not found") + + metadata = page.metadata.get(item) + if metadata is None: + raise HTTPException(status_code=404, detail="File not found") + + cover_images = json.loads(metadata.get('ssmd_cover_images', {})) + image = cover_images[index] if index < len(cover_images) else None + if not image: + raise HTTPException(status_code=404, detail="File not found") + + try: + image = Image.open(BytesIO(b64decode(image))) + buffer = BytesIO() + image.save(buffer, format=image.format) + return Response(content=buffer.getvalue(), media_type=image.get_format_mimetype()) + except Exception as err: + raise ValueError(f"File cannot be fetched: {item}. Failed to load cover image.") from err + + def get_metadata(page: str = "", item: str = ""): from starlette.responses import JSONResponse @@ -119,6 +147,8 @@ def get_metadata(page: str = "", item: str = ""): if metadata is None: return JSONResponse({}) + metadata = {i:metadata[i] for i in metadata if i != 'ssmd_cover_images'} # those are cover images, and they are too big to display in UI as text + return JSONResponse({"metadata": json.dumps(metadata, indent=4, ensure_ascii=False)}) @@ -142,6 +172,7 @@ def get_single_card(page: str = "", tabname: str = "", name: str = ""): def add_pages_to_demo(app): app.add_api_route("/sd_extra_networks/thumb", fetch_file, methods=["GET"]) + app.add_api_route("/sd_extra_networks/cover-images", fetch_cover_images, methods=["GET"]) app.add_api_route("/sd_extra_networks/metadata", get_metadata, methods=["GET"]) app.add_api_route("/sd_extra_networks/get-single-card", get_single_card, methods=["GET"]) @@ -627,6 +658,17 @@ class ExtraNetworksPage: return None + def find_embedded_preview(self, path, name, metadata): + """ + Find if embedded preview exists in safetensors metadata and return endpoint for it. + """ + + file = f"{path}.safetensors" + if self.lister.exists(file) and 'ssmd_cover_images' in metadata and len(list(filter(None, json.loads(metadata['ssmd_cover_images'])))) > 0: + return f"./sd_extra_networks/cover-images?page={self.extra_networks_tabname}&item={name}" + + return None + def find_description(self, path): """ Find and read a description file for a given path (without extension).