From 27a3fa8f516cebeec3a3a8d0f25cb0e321851807 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 12 Feb 2023 09:51:35 -0600 Subject: [PATCH] feat: show model in image card, use labels for model and scheduler (#104) --- api/onnx_web/diffusion/load.py | 38 +++++++++++++++++++++++++++++++- api/onnx_web/output.py | 5 +++++ api/onnx_web/serve.py | 33 +++------------------------ gui/src/client.ts | 2 +- gui/src/components/ImageCard.tsx | 11 ++++++--- gui/src/strings.ts | 2 +- 6 files changed, 55 insertions(+), 36 deletions(-) diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py index 9256f1aa..476d56e1 100644 --- a/api/onnx_web/diffusion/load.py +++ b/api/onnx_web/diffusion/load.py @@ -2,7 +2,21 @@ from logging import getLogger from typing import Any, Optional, Tuple import numpy as np -from diffusers import DiffusionPipeline +from diffusers import ( + DDIMScheduler, + DDPMScheduler, + DiffusionPipeline, + DPMSolverMultistepScheduler, + DPMSolverSinglestepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + KarrasVeScheduler, + KDPM2AncestralDiscreteScheduler, + KDPM2DiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) from ..params import DeviceParams, Size from ..utils import run_gc @@ -22,6 +36,28 @@ last_pipeline_scheduler: Any = None latent_channels = 4 latent_factor = 8 +pipeline_schedulers = { + "ddim": DDIMScheduler, + "ddpm": DDPMScheduler, + "dpm-multi": DPMSolverMultistepScheduler, + "dpm-single": DPMSolverSinglestepScheduler, + "euler": EulerDiscreteScheduler, + "euler-a": EulerAncestralDiscreteScheduler, + "heun": HeunDiscreteScheduler, + "k-dpm-2-a": KDPM2AncestralDiscreteScheduler, + "k-dpm-2": KDPM2DiscreteScheduler, + "karras-ve": KarrasVeScheduler, + "lms-discrete": LMSDiscreteScheduler, + "pndm": PNDMScheduler, +} + +def get_scheduler_name(scheduler: Any) -> Optional[str]: + for k, v in pipeline_schedulers.items(): + if scheduler == v or scheduler == v.__name__: + return k + + return None + def get_latents_from_seed(seed: int, size: Size, batch: int = 1) -> np.ndarray: """ diff --git a/api/onnx_web/output.py b/api/onnx_web/output.py index 0b65641c..645b6ad3 100644 --- a/api/onnx_web/output.py +++ b/api/onnx_web/output.py @@ -1,12 +1,14 @@ from hashlib import sha256 from json import dumps from logging import getLogger +from os import path from struct import pack from time import time from typing import Any, Optional, Tuple from PIL import Image +from .diffusion.load import get_scheduler_name from .params import Border, ImageParams, Param, Size, UpscaleParams from .utils import ServerContext, base_join @@ -38,6 +40,9 @@ def json_params( "params": params.tojson(), } + json["params"]["model"] = path.basename(params.model) + json["params"]["scheduler"] = get_scheduler_name(params.scheduler) + if upscale is not None and border is not None: size = upscale.resize(size.add_border(border)) diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index e2a05d2e..c92b7629 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -9,28 +9,12 @@ from typing import Dict, List, Tuple, Union import numpy as np import torch import yaml -from diffusers import ( - DDIMScheduler, - DDPMScheduler, - DPMSolverMultistepScheduler, - DPMSolverSinglestepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - KarrasVeScheduler, - KDPM2AncestralDiscreteScheduler, - KDPM2DiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) from flask import Flask, jsonify, make_response, request, send_from_directory, url_for from flask_cors import CORS from jsonschema import validate from onnxruntime import get_available_providers from PIL import Image -from onnx_web.hacks import apply_patches - from .chain import ( ChainPipeline, blend_img2img, @@ -48,12 +32,14 @@ from .chain import ( upscale_stable_diffusion, ) from .device_pool import DevicePoolExecutor +from .diffusion.load import pipeline_schedulers from .diffusion.run import ( run_img2img_pipeline, run_inpaint_pipeline, run_txt2img_pipeline, run_upscale_pipeline, ) +from .hacks import apply_patches from .image import ( # mask filters; noise sources mask_filter_gaussian_multiply, mask_filter_gaussian_screen, @@ -99,20 +85,7 @@ platform_providers = { "directml": "DmlExecutionProvider", "rocm": "ROCMExecutionProvider", } -pipeline_schedulers = { - "ddim": DDIMScheduler, - "ddpm": DDPMScheduler, - "dpm-multi": DPMSolverMultistepScheduler, - "dpm-single": DPMSolverSinglestepScheduler, - "euler": EulerDiscreteScheduler, - "euler-a": EulerAncestralDiscreteScheduler, - "heun": HeunDiscreteScheduler, - "k-dpm-2-a": KDPM2AncestralDiscreteScheduler, - "k-dpm-2": KDPM2DiscreteScheduler, - "karras-ve": KarrasVeScheduler, - "lms-discrete": LMSDiscreteScheduler, - "pndm": PNDMScheduler, -} + noise_sources = { "fill-edge": noise_source_fill_edge, "fill-mask": noise_source_fill_mask, diff --git a/gui/src/client.ts b/gui/src/client.ts index b8f9bb56..e68ee42c 100644 --- a/gui/src/client.ts +++ b/gui/src/client.ts @@ -137,7 +137,7 @@ export interface ImageResponse { key: string; url: string; }; - params: Required; + params: Required & Required; size: { width: number; height: number; diff --git a/gui/src/components/ImageCard.tsx b/gui/src/components/ImageCard.tsx index a2dafcea..daa95744 100644 --- a/gui/src/components/ImageCard.tsx +++ b/gui/src/components/ImageCard.tsx @@ -1,4 +1,4 @@ -import { doesExist, mustExist } from '@apextoaster/js-utils'; +import { doesExist, mustDefault, mustExist } from '@apextoaster/js-utils'; import { Brush, ContentCopy, Delete, Download } from '@mui/icons-material'; import { Box, Card, CardContent, CardMedia, Grid, IconButton, Paper, Tooltip } from '@mui/material'; import * as React from 'react'; @@ -8,6 +8,7 @@ import { useStore } from 'zustand'; import { ImageResponse } from '../client.js'; import { ConfigContext, StateContext } from '../state.js'; +import { MODEL_LABELS, SCHEDULER_LABELS } from '../strings.js'; export interface ImageCardProps { value: ImageResponse; @@ -64,6 +65,9 @@ export function ImageCard(props: ImageCardProps) { window.open(output.url, '_blank'); } + const model = mustDefault(MODEL_LABELS[params.model], params.model); + const scheduler = mustDefault(SCHEDULER_LABELS[params.scheduler], params.scheduler); + return + Model: {model} + Scheduler: {scheduler} + Seed: {params.seed} CFG: {params.cfg} Steps: {params.steps} Size: {size.width}x{size.height} - Seed: {params.seed} - Scheduler: {params.scheduler} {params.prompt} diff --git a/gui/src/strings.ts b/gui/src/strings.ts index 79c7e689..b58f525b 100644 --- a/gui/src/strings.ts +++ b/gui/src/strings.ts @@ -1,5 +1,5 @@ // TODO: set up i18next -export const MODEL_LABELS = { +export const MODEL_LABELS: Record = { 'stable-diffusion-onnx-v1-4': 'Stable Diffusion v1.4', 'stable-diffusion-onnx-v1-5': 'Stable Diffusion v1.5', 'stable-diffusion-onnx-v1-inpainting': 'SD Inpainting v1',