add upscale support

This commit is contained in:
jtydhr88
2023-05-22 16:50:11 -04:00
parent 7d73a79671
commit f87e36598b
2 changed files with 134 additions and 45 deletions

View File

@@ -24,9 +24,9 @@ Still, currently this plugin is a basic implement for webui, and only support a
- [x] webui status could replace images generated status
- [x] load existing images (need an extension on webui: [sd-webui-StableStudio](https://github.com/jtydhr88/sd-webui-StableStudio))
- [x] Need to think about how to deal with extensions ecosystem in webui (made a start with [sd-webui-StableStudio](https://github.com/jtydhr88/sd-webui-StableStudio))
- [ ] upscale (working)
- [x] upscale
- [ ] test and bugs fix (working)
- [ ] Lora support
- [ ] plugin could use path in settings along with a field for extra cli flags to launch webui on startup
- [ ] some bugs fix
- [ ] MacOS and Linux support (since I dont have environment with MacOS/Linux, may need someone help with this)
- [ ] many other features from webui...

View File

@@ -15,11 +15,55 @@ function blobToBase64(blob: Blob): Promise<string> {
});
}
const webuiUpscalers = [
{
label: "None",
value: "None"
},
{
label: "Lanczos",
value: "Lanczos"
},
{
label: "Nearest",
value: "Nearest"
},
{
label: "ESRGAN_4x",
value: "ESRGAN_4x"
},
{
label: "LDSR",
value: "LDSR"
},
{
label: "R-ESRGAN 4x+",
value: "R-ESRGAN 4x+"
},
{
label: "R-ESRGAN 4x+ Anime6B",
value: "R-ESRGAN 4x+ Anime6B"
},
{
label: "ScuNET GAN",
value: "ScuNET GAN"
},
{
label: "ScuNET PSNR",
value: "ScuNET PSNR"
},
{
label: "SwinIR_4x",
value: "SwinIR_4x"
}
];
const getStableDiffusionDefaultCount = () => 4;
export const createPlugin = StableStudio.createPlugin<{
imagesGeneratedSoFar: number;
settings: {
webuiHostUrl: StableStudio.PluginSettingString;
upscaler1: StableStudio.PluginSettingString;
};
}>(({set, get}) => {
const webuiLoad = (webuiHostUrl?: string): Pick<
@@ -51,17 +95,24 @@ export const createPlugin = StableStudio.createPlugin<{
const webUIOptions = await optionsResponse.json();
console.log(webUIOptions);
const model = options?.input?.model;
let url = webuiHostUrl + '/sdapi/v1/txt2img';
let isUpscale = false;
if (options?.input?.initialImage) {
url = webuiHostUrl + '/sdapi/v1/img2img';
// use this condition to assume the user wants to upscale
if ((options?.input?.initialImage?.weight === 1) && (model === "esrgan-v1-x2plus")) {
isUpscale = true;
url = webuiHostUrl + '/sdapi/v1/extra-single-image';
}
}
const model = options?.input?.model;
if (model && model !== webUIOptions["sd_model_checkpoint"]) {
if (model && model !== webUIOptions["sd_model_checkpoint"] && !isUpscale) {
console.log("applying model");
localStorage.setItem("model", model)
@@ -78,7 +129,9 @@ export const createPlugin = StableStudio.createPlugin<{
body: JSON.stringify(modelData),
});
console.log(modelResponse);
if (modelResponse.ok) {
console.log("applied model");
}
}
const sampler = options?.input?.sampler?.name;
@@ -107,22 +160,36 @@ export const createPlugin = StableStudio.createPlugin<{
let data: dataObject = {};
data["seed"] = seed;
data["sampler_name"] = options?.input?.sampler?.name ?? "";
data["sampler_index"] = options?.input?.sampler?.name ?? "";
data["prompt"] = prompt;
data["negative_prompt"] = negativePrompt;
data["steps"] = options?.input?.steps ?? 20;
data["batch_size"] = options?.count ?? getStableDiffusionDefaultCount();
data["width"] = options?.input?.width ?? 512;
localStorage.setItem("width", data["width"]);
if (!isUpscale) {
data["seed"] = seed;
data["sampler_name"] = options?.input?.sampler?.name ?? "";
data["sampler_index"] = options?.input?.sampler?.name ?? "";
data["prompt"] = prompt;
data["negative_prompt"] = negativePrompt;
data["steps"] = options?.input?.steps ?? 20;
data["batch_size"] = options?.count ?? getStableDiffusionDefaultCount();
data["height"] = options?.input?.height ?? 512;
data["width"] = options?.input?.width ?? 512;
localStorage.setItem("height", data["height"]);
localStorage.setItem("width", data["width"]);
data["save_images"] = true;
data["height"] = options?.input?.height ?? 512;
localStorage.setItem("height", data["height"]);
data["save_images"] = true;
} else {
data["upscaling_resize_w"] = options?.input?.width ?? 512;
data["upscaling_resize_h"] = options?.input?.height ?? 512;
data["upscaler_1"] = get().settings.upscaler1.value;
}
if (options?.input?.initialImage?.weight && !isUpscale) {
data["denoising_strength"] = 1 - options.input.initialImage.weight;
}
if (options?.input?.cfgScale) {
data["cfg_scale"] = options?.input?.cfgScale;
@@ -130,14 +197,14 @@ export const createPlugin = StableStudio.createPlugin<{
localStorage.setItem("cfgScale", data["cfg_scale"])
}
if (options?.input?.initialImage?.weight) {
data["denoising_strength"] = 1 - options.input.initialImage.weight;
}
if (options?.input?.initialImage?.blob) {
const initImgB64 = await blobToBase64(options?.input?.initialImage?.blob);
data["init_images"] = [initImgB64.split(",")[1]];
if (isUpscale) {
data["image"] = initImgB64.split(",")[1];
} else {
data["init_images"] = [initImgB64.split(",")[1]];
}
}
if (options?.input?.maskImage?.blob) {
@@ -162,35 +229,49 @@ export const createPlugin = StableStudio.createPlugin<{
const responseData = await response.json();
const images = [];
console.log(responseData)
const images = [];
const createdAt = new Date();
console.log(responseData.images);
const generatedImagesLength = responseData.images.length;
let startIndex = 0;
if (generatedImagesLength > data["batch_size"]) {
startIndex = 1;
}
for (let i = startIndex; i < responseData.images.length; i++) {
const blob = await base64ToBlob(responseData.images[i], 'image/jpeg');
if (isUpscale) {
const blob = await base64ToBlob(responseData.image, 'image/jpeg');
const image = {
id: `${Math.random() * 10000000}`,
createdAt,
blob
createdAt: createdAt,
blob: blob,
input: {
model: localStorage.getItem("model")
}
}
images.push(image)
}
images.push(image);
} else {
const generatedImagesLength = responseData.images.length;
set(({imagesGeneratedSoFar}) => ({
imagesGeneratedSoFar: imagesGeneratedSoFar + data["batch_size"],
}));
let startIndex = 0;
if (generatedImagesLength > data["batch_size"]) {
startIndex = 1;
}
for (let i = startIndex; i < responseData.images.length; i++) {
const blob = await base64ToBlob(responseData.images[i], 'image/jpeg');
const image = {
id: `${Math.random() * 10000000}`,
createdAt,
blob
}
images.push(image)
}
set(({imagesGeneratedSoFar}) => ({
imagesGeneratedSoFar: imagesGeneratedSoFar + data["batch_size"],
}));
}
return {
id: `${Math.random() * 10000000}`,
@@ -356,10 +437,16 @@ export const createPlugin = StableStudio.createPlugin<{
type: "string",
title: "Webui Host URL",
description:
"put your webui url here, the default value is http://127.0.0.1:7861",
"put your webui api url here, the default value is http://127.0.0.1:7861",
placeholder: "http://127.0.0.1:7861",
value: localStorage.getItem("webui-host-url") ?? "",
},
upscaler1: {
type: "string",
title: "Upscaler 1",
options: webuiUpscalers,
value: localStorage.getItem("upscaler1") ?? "None",
}
},
setSetting: (key, value) => {
@@ -373,6 +460,8 @@ export const createPlugin = StableStudio.createPlugin<{
if (key === "webuiHostUrl" && typeof value === "string") {
localStorage.setItem("webui-host-url", value);
set((plugin) => ({...plugin, ...webuiLoad(value)}));
} else if (key === "upscaler1" && typeof value === "string") {
localStorage.setItem("upscaler1", value);
}
},