Skip to content

Commit f87e365

Browse files
jtydhr88jtydhr88
committed
add upscale support
1 parent 7d73a79 commit f87e365

File tree

2 files changed

+132
-43
lines changed

2 files changed

+132
-43
lines changed

‎packages/stablestudio-plugin-webui/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,9 @@ Still, currently this plugin is a basic implement for webui, and only support a
2424
- [x] webui status could replace images generated status
2525
- [x] load existing images (need an extension on webui: [sd-webui-StableStudio](https://github.com/jtydhr88/sd-webui-StableStudio))
2626
- [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))
27-
- [ ] upscale (working)
27+
- [x] upscale
28+
- [ ] test and bugs fix (working)
2829
- [ ] Lora support
2930
- [ ] plugin could use path in settings along with a field for extra cli flags to launch webui on startup
30-
- [ ] some bugs fix
3131
- [ ] MacOS and Linux support (since I don’t have environment with MacOS/Linux, may need someone help with this)
3232
- [ ] many other features from webui...

‎packages/stablestudio-plugin-webui/src/index.ts‎

Lines changed: 130 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,55 @@ function blobToBase64(blob: Blob): Promise<string> {
1515
});
1616
}
1717

18+
const webuiUpscalers = [
19+
{
20+
label: "None",
21+
value: "None"
22+
},
23+
{
24+
label: "Lanczos",
25+
value: "Lanczos"
26+
},
27+
{
28+
label: "Nearest",
29+
value: "Nearest"
30+
},
31+
{
32+
label: "ESRGAN_4x",
33+
value: "ESRGAN_4x"
34+
},
35+
{
36+
label: "LDSR",
37+
value: "LDSR"
38+
},
39+
{
40+
label: "R-ESRGAN 4x+",
41+
value: "R-ESRGAN 4x+"
42+
},
43+
{
44+
label: "R-ESRGAN 4x+ Anime6B",
45+
value: "R-ESRGAN 4x+ Anime6B"
46+
},
47+
{
48+
label: "ScuNET GAN",
49+
value: "ScuNET GAN"
50+
},
51+
{
52+
label: "ScuNET PSNR",
53+
value: "ScuNET PSNR"
54+
},
55+
{
56+
label: "SwinIR_4x",
57+
value: "SwinIR_4x"
58+
}
59+
];
60+
1861
const getStableDiffusionDefaultCount = () => 4;
1962
export const createPlugin = StableStudio.createPlugin<{
2063
imagesGeneratedSoFar: number;
2164
settings: {
2265
webuiHostUrl: StableStudio.PluginSettingString;
66+
upscaler1: StableStudio.PluginSettingString;
2367
};
2468
}>(({set, get}) => {
2569
const webuiLoad = (webuiHostUrl?: string): Pick<
@@ -51,17 +95,24 @@ export const createPlugin = StableStudio.createPlugin<{
5195

5296
const webUIOptions = await optionsResponse.json();
5397

54-
console.log(webUIOptions);
98+
const model = options?.input?.model;
5599

56100
let url = webuiHostUrl + '/sdapi/v1/txt2img';
57101

102+
let isUpscale = false;
103+
58104
if (options?.input?.initialImage) {
59105
url = webuiHostUrl + '/sdapi/v1/img2img';
60-
}
61106

62-
const model = options?.input?.model;
107+
// use this condition to assume the user wants to upscale
108+
if ((options?.input?.initialImage?.weight === 1) && (model === "esrgan-v1-x2plus")) {
109+
isUpscale = true;
110+
111+
url = webuiHostUrl + '/sdapi/v1/extra-single-image';
112+
}
113+
}
63114

64-
if (model && model !== webUIOptions["sd_model_checkpoint"]) {
115+
if (model && model !== webUIOptions["sd_model_checkpoint"] && !isUpscale) {
65116
console.log("applying model");
66117

67118
localStorage.setItem("model", model)
@@ -78,7 +129,9 @@ export const createPlugin = StableStudio.createPlugin<{
78129
body: JSON.stringify(modelData),
79130
});
80131

81-
console.log(modelResponse);
132+
if (modelResponse.ok) {
133+
console.log("applied model");
134+
}
82135
}
83136

84137
const sampler = options?.input?.sampler?.name;
@@ -107,37 +160,51 @@ export const createPlugin = StableStudio.createPlugin<{
107160

108161
let data: dataObject = {};
109162

110-
data["seed"] = seed;
111-
data["sampler_name"] = options?.input?.sampler?.name ?? "";
112-
data["sampler_index"] = options?.input?.sampler?.name ?? "";
113-
data["prompt"] = prompt;
114-
data["negative_prompt"] = negativePrompt;
115-
data["steps"] = options?.input?.steps ?? 20;
116-
data["batch_size"] = options?.count ?? getStableDiffusionDefaultCount();
117-
data["width"] = options?.input?.width ?? 512;
118163

119-
localStorage.setItem("width", data["width"]);
164+
if (!isUpscale) {
165+
data["seed"] = seed;
166+
data["sampler_name"] = options?.input?.sampler?.name ?? "";
167+
data["sampler_index"] = options?.input?.sampler?.name ?? "";
168+
data["prompt"] = prompt;
169+
data["negative_prompt"] = negativePrompt;
170+
data["steps"] = options?.input?.steps ?? 20;
171+
data["batch_size"] = options?.count ?? getStableDiffusionDefaultCount();
120172

121-
data["height"] = options?.input?.height ?? 512;
173+
data["width"] = options?.input?.width ?? 512;
122174

123-
localStorage.setItem("height", data["height"]);
175+
localStorage.setItem("width", data["width"]);
124176

125-
data["save_images"] = true;
177+
data["height"] = options?.input?.height ?? 512;
126178

127-
if (options?.input?.cfgScale) {
128-
data["cfg_scale"] = options?.input?.cfgScale;
179+
localStorage.setItem("height", data["height"]);
129180

130-
localStorage.setItem("cfgScale", data["cfg_scale"])
181+
data["save_images"] = true;
182+
} else {
183+
data["upscaling_resize_w"] = options?.input?.width ?? 512;
184+
185+
data["upscaling_resize_h"] = options?.input?.height ?? 512;
186+
187+
data["upscaler_1"] = get().settings.upscaler1.value;
131188
}
132189

133-
if (options?.input?.initialImage?.weight) {
190+
if (options?.input?.initialImage?.weight && !isUpscale) {
134191
data["denoising_strength"] = 1 - options.input.initialImage.weight;
135192
}
136193

194+
if (options?.input?.cfgScale) {
195+
data["cfg_scale"] = options?.input?.cfgScale;
196+
197+
localStorage.setItem("cfgScale", data["cfg_scale"])
198+
}
199+
137200
if (options?.input?.initialImage?.blob) {
138201
const initImgB64 = await blobToBase64(options?.input?.initialImage?.blob);
139202

140-
data["init_images"] = [initImgB64.split(",")[1]];
203+
if (isUpscale) {
204+
data["image"] = initImgB64.split(",")[1];
205+
} else {
206+
data["init_images"] = [initImgB64.split(",")[1]];
207+
}
141208
}
142209

143210
if (options?.input?.maskImage?.blob) {
@@ -162,36 +229,50 @@ export const createPlugin = StableStudio.createPlugin<{
162229

163230
const responseData = await response.json();
164231

165-
const images = [];
232+
console.log(responseData)
166233

234+
const images = [];
167235
const createdAt = new Date();
168236

169-
console.log(responseData.images);
237+
if (isUpscale) {
238+
const blob = await base64ToBlob(responseData.image, 'image/jpeg');
239+
240+
const image = {
241+
id: `${Math.random() * 10000000}`,
242+
createdAt: createdAt,
243+
blob: blob,
244+
input: {
245+
model: localStorage.getItem("model")
246+
}
247+
}
248+
249+
images.push(image);
250+
} else {
251+
const generatedImagesLength = responseData.images.length;
170252

171-
const generatedImagesLength = responseData.images.length;
253+
let startIndex = 0;
172254

173-
let startIndex = 0;
255+
if (generatedImagesLength > data["batch_size"]) {
256+
startIndex = 1;
257+
}
174258

175-
if (generatedImagesLength > data["batch_size"]) {
176-
startIndex = 1;
177-
}
259+
for (let i = startIndex; i < responseData.images.length; i++) {
260+
const blob = await base64ToBlob(responseData.images[i], 'image/jpeg');
178261

179-
for (let i = startIndex; i < responseData.images.length; i++) {
180-
const blob = await base64ToBlob(responseData.images[i], 'image/jpeg');
262+
const image = {
263+
id: `${Math.random() * 10000000}`,
264+
createdAt,
265+
blob
266+
}
181267

182-
const image = {
183-
id: `${Math.random() * 10000000}`,
184-
createdAt,
185-
blob
268+
images.push(image)
186269
}
187270

188-
images.push(image)
271+
set(({imagesGeneratedSoFar}) => ({
272+
imagesGeneratedSoFar: imagesGeneratedSoFar + data["batch_size"],
273+
}));
189274
}
190275

191-
set(({imagesGeneratedSoFar}) => ({
192-
imagesGeneratedSoFar: imagesGeneratedSoFar + data["batch_size"],
193-
}));
194-
195276
return {
196277
id: `${Math.random() * 10000000}`,
197278
images: images
@@ -356,10 +437,16 @@ export const createPlugin = StableStudio.createPlugin<{
356437
type: "string",
357438
title: "Webui Host URL",
358439
description:
359-
"put your webui url here, the default value is http://127.0.0.1:7861",
440+
"put your webui api url here, the default value is http://127.0.0.1:7861",
360441
placeholder: "http://127.0.0.1:7861",
361442
value: localStorage.getItem("webui-host-url") ?? "",
362443
},
444+
upscaler1: {
445+
type: "string",
446+
title: "Upscaler 1",
447+
options: webuiUpscalers,
448+
value: localStorage.getItem("upscaler1") ?? "None",
449+
}
363450
},
364451

365452
setSetting: (key, value) => {
@@ -373,6 +460,8 @@ export const createPlugin = StableStudio.createPlugin<{
373460
if (key === "webuiHostUrl" && typeof value === "string") {
374461
localStorage.setItem("webui-host-url", value);
375462
set((plugin) => ({...plugin, ...webuiLoad(value)}));
463+
} else if (key === "upscaler1" && typeof value === "string") {
464+
localStorage.setItem("upscaler1", value);
376465
}
377466
},
378467

0 commit comments

Comments
 (0)