diff --git a/api/onnx_web/chain/result.py b/api/onnx_web/chain/result.py index 7166f397..04d02f3c 100644 --- a/api/onnx_web/chain/result.py +++ b/api/onnx_web/chain/result.py @@ -1,13 +1,14 @@ from json import dumps from logging import getLogger from os import path -from typing import Any, List, Optional +from typing import Any, List, Optional, Tuple import numpy as np from PIL import Image from ..convert.utils import resolve_tensor from ..params import Border, HighresParams, ImageParams, Size, UpscaleParams +from ..server.context import ServerContext from ..server.load import get_extra_hashes from ..utils import hash_file @@ -55,8 +56,8 @@ class ImageMetadata: self.loras = loras self.models = models - def to_auto1111(self, server, outputs) -> str: - model_name = path.basename(path.normpath(self.params.model)) + def get_model_hash(self, model: Optional[str] = None) -> Tuple[str, str]: + model_name = path.basename(path.normpath(model or self.params.model)) logger.debug("getting model hash for %s", model_name) model_hash = get_extra_hashes().get(model_name, None) @@ -66,7 +67,10 @@ class ImageMetadata: with open(model_hash_path, "r") as f: model_hash = f.readline().rstrip(",. \n\t\r") - model_hash = model_hash or "unknown" + return (model_name, model_hash or "unknown") + + def to_exif(self, server) -> str: + model_name, model_hash = self.get_model_hash() hash_map = { model_name: model_hash, } @@ -112,15 +116,17 @@ class ImageMetadata: f"Hashes: {dumps(hash_map)}" ) - def tojson(self, server, outputs): + def tojson(self, server: ServerContext, output: List[str]): json = { "input_size": self.size.tojson(), - "outputs": outputs, + "outputs": output, "params": self.params.tojson(), - "inversions": {}, - "loras": {}, + "inversions": [], + "loras": [], + "models": [], } + # fix up some fields json["params"]["model"] = path.basename(self.params.model) json["params"]["scheduler"] = self.params.scheduler # TODO: why tho? @@ -145,14 +151,21 @@ class ImageMetadata: hash = hash_file( resolve_tensor(path.join(server.model_path, "inversion", name)) ).upper() - json["inversions"][name] = {"weight": weight, "hash": hash} + json["inversions"].append( + {"name": name, "weight": weight, "hash": hash} + ) if self.loras is not None: for name, weight in self.loras: hash = hash_file( resolve_tensor(path.join(server.model_path, "lora", name)) ).upper() - json["loras"][name] = {"weight": weight, "hash": hash} + json["loras"].append({"name": name, "weight": weight, "hash": hash}) + + if self.models is not None: + for name, weight in self.models: + name, hash = self.get_model_hash() + json["models"].append({"name": name, "weight": weight, "hash": hash}) return json diff --git a/api/onnx_web/output.py b/api/onnx_web/output.py index 9173d168..64c2f8ad 100644 --- a/api/onnx_web/output.py +++ b/api/onnx_web/output.py @@ -98,7 +98,7 @@ def save_image( exif.add_text("model", server.server_version) exif.add_text( "parameters", - metadata.to_auto1111(server, [output]), + metadata.to_exif(server, [output]), ) image.save(path, format=server.image_format, pnginfo=exif) @@ -111,7 +111,7 @@ def save_image( encoding="unicode", ), ExifIFD.UserComment: UserComment.dump( - metadata.to_auto1111(server, [output]), + metadata.to_exif(server, [output]), encoding="unicode", ), ImageIFD.Make: "onnx-web", diff --git a/gui/src/components/card/ImageCard.tsx b/gui/src/components/card/ImageCard.tsx index 75a2ac0c..249f6c6f 100644 --- a/gui/src/components/card/ImageCard.tsx +++ b/gui/src/components/card/ImageCard.tsx @@ -12,7 +12,6 @@ import { ClientContext, ConfigContext, OnnxState, StateContext } from '../../sta import { range, visibleIndex } from '../../utils.js'; import { BLEND_SOURCES } from '../../constants.js'; import { JobResponse, SuccessJobResponse } from '../../types/api-v2.js'; -import { getApiRoot } from '../../config.js'; export interface ImageCardProps { image: SuccessJobResponse; @@ -110,8 +109,8 @@ export function ImageCard(props: ImageCardProps) { const url = useMemo(() => client.outputURL(image, index), [image, index]); - const model = getLabel('model', metadata[index].model); - const scheduler = getLabel('scheduler', metadata[index].scheduler); + const model = getLabel('model', metadata[index].models[0].name); + const scheduler = getLabel('scheduler', metadata[index].params.scheduler); return { - if (ready.status === 'success' && getReady()) { + if (ready.status === 'success') { setReady(ready.data[0]); } - }, [ready.status, getReady(), getProgress(ready.data)]); + }, [ready.status, getStatus(ready.data), getProgress(ready.data)]); return @@ -134,3 +130,11 @@ export function getTotal(data: Maybe>) { return 0; } + +function getStatus(data: Maybe>) { + if (doesExist(data)) { + return data[0].status; + } + + return JobStatus.PENDING; +} diff --git a/gui/src/types/api-v2.ts b/gui/src/types/api-v2.ts index 1135a6a6..afafcc9e 100644 --- a/gui/src/types/api-v2.ts +++ b/gui/src/types/api-v2.ts @@ -19,16 +19,18 @@ export interface NetworkMetadata { export interface ImageMetadata { input_size: Size; - size: Size; outputs: Array; params: TParams; + inversions: Array; loras: Array; - model: string; - scheduler: string; - border: unknown; + models: Array; + + border: unknown; // TODO: type highres: HighresParams; upscale: UpscaleParams; + size: Size; + type: TType; }