2023-01-13 20:39:07 +00:00
|
|
|
/* eslint-disable no-null/no-null */
|
2023-01-13 00:10:46 +00:00
|
|
|
import { Maybe } from '@apextoaster/js-utils';
|
|
|
|
import { createContext } from 'react';
|
|
|
|
import { StateCreator, StoreApi } from 'zustand';
|
|
|
|
|
|
|
|
import {
|
|
|
|
ApiClient,
|
|
|
|
ApiResponse,
|
|
|
|
BaseImgParams,
|
2023-01-15 00:07:05 +00:00
|
|
|
BrushParams,
|
2023-01-13 00:10:46 +00:00
|
|
|
Img2ImgParams,
|
|
|
|
InpaintParams,
|
2023-01-17 02:11:10 +00:00
|
|
|
ModelParams,
|
2023-01-15 00:07:05 +00:00
|
|
|
OutpaintPixels,
|
2023-01-13 00:10:46 +00:00
|
|
|
paramsFromConfig,
|
|
|
|
Txt2ImgParams,
|
2023-01-16 20:05:04 +00:00
|
|
|
UpscaleParams,
|
2023-01-16 20:10:15 +00:00
|
|
|
} from './client.js';
|
2023-01-13 20:39:07 +00:00
|
|
|
import { ConfigFiles, ConfigParams, ConfigState } from './config.js';
|
2023-01-12 06:10:57 +00:00
|
|
|
|
2023-01-13 20:39:07 +00:00
|
|
|
type TabState<TabParams extends BaseImgParams> = ConfigFiles<Required<TabParams>> & ConfigState<Required<TabParams>>;
|
2023-01-12 06:10:57 +00:00
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
interface Txt2ImgSlice {
|
|
|
|
txt2img: TabState<Txt2ImgParams>;
|
|
|
|
|
|
|
|
setTxt2Img(params: Partial<Txt2ImgParams>): void;
|
|
|
|
resetTxt2Img(): void;
|
2023-01-12 06:10:57 +00:00
|
|
|
}
|
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
interface Img2ImgSlice {
|
|
|
|
img2img: TabState<Img2ImgParams>;
|
2023-01-12 06:10:57 +00:00
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
setImg2Img(params: Partial<Img2ImgParams>): void;
|
|
|
|
resetImg2Img(): void;
|
|
|
|
}
|
2023-01-12 06:10:57 +00:00
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
interface InpaintSlice {
|
|
|
|
inpaint: TabState<InpaintParams>;
|
2023-01-12 06:10:57 +00:00
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
setInpaint(params: Partial<InpaintParams>): void;
|
|
|
|
resetInpaint(): void;
|
|
|
|
}
|
|
|
|
|
|
|
|
interface HistorySlice {
|
|
|
|
history: Array<ApiResponse>;
|
|
|
|
limit: number;
|
2023-01-13 03:12:20 +00:00
|
|
|
loading: Maybe<ApiResponse>;
|
2023-01-12 23:49:51 +00:00
|
|
|
|
|
|
|
pushHistory(image: ApiResponse): void;
|
|
|
|
removeHistory(image: ApiResponse): void;
|
|
|
|
setLimit(limit: number): void;
|
2023-01-13 03:12:20 +00:00
|
|
|
setLoading(image: Maybe<ApiResponse>): void;
|
2023-01-12 23:49:51 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
interface DefaultSlice {
|
|
|
|
defaults: TabState<BaseImgParams>;
|
|
|
|
|
|
|
|
setDefaults(param: Partial<BaseImgParams>): void;
|
2023-01-12 06:10:57 +00:00
|
|
|
}
|
|
|
|
|
2023-01-15 00:07:05 +00:00
|
|
|
interface OutpaintSlice {
|
|
|
|
outpaint: OutpaintPixels;
|
|
|
|
|
|
|
|
setOutpaint(pixels: Partial<OutpaintPixels>): void;
|
|
|
|
}
|
|
|
|
|
|
|
|
interface BrushSlice {
|
|
|
|
brush: BrushParams;
|
|
|
|
|
|
|
|
setBrush(brush: Partial<BrushParams>): void;
|
|
|
|
}
|
|
|
|
|
2023-01-16 20:05:04 +00:00
|
|
|
interface UpscaleSlice {
|
|
|
|
upscale: UpscaleParams;
|
|
|
|
|
|
|
|
setUpscale(upscale: Partial<UpscaleParams>): void;
|
|
|
|
}
|
|
|
|
|
2023-01-17 02:11:10 +00:00
|
|
|
interface ModelSlice {
|
|
|
|
model: ModelParams;
|
|
|
|
|
|
|
|
setModel(model: Partial<ModelParams>): void;
|
|
|
|
}
|
|
|
|
|
2023-01-16 20:05:04 +00:00
|
|
|
export type OnnxState
|
|
|
|
= BrushSlice
|
|
|
|
& DefaultSlice
|
|
|
|
& HistorySlice
|
|
|
|
& Img2ImgSlice
|
|
|
|
& InpaintSlice
|
2023-01-17 02:11:10 +00:00
|
|
|
& ModelSlice
|
2023-01-16 20:05:04 +00:00
|
|
|
& OutpaintSlice
|
|
|
|
& Txt2ImgSlice
|
|
|
|
& UpscaleSlice;
|
2023-01-12 23:49:51 +00:00
|
|
|
|
|
|
|
export function createStateSlices(base: ConfigParams) {
|
|
|
|
const defaults = paramsFromConfig(base);
|
|
|
|
|
|
|
|
const createTxt2ImgSlice: StateCreator<OnnxState, [], [], Txt2ImgSlice> = (set) => ({
|
|
|
|
txt2img: {
|
|
|
|
...defaults,
|
|
|
|
width: base.width.default,
|
|
|
|
height: base.height.default,
|
|
|
|
},
|
|
|
|
setTxt2Img(params) {
|
|
|
|
set((prev) => ({
|
|
|
|
txt2img: {
|
|
|
|
...prev.txt2img,
|
|
|
|
...params,
|
|
|
|
},
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
resetTxt2Img() {
|
|
|
|
set({
|
|
|
|
txt2img: {
|
|
|
|
...defaults,
|
|
|
|
width: base.width.default,
|
|
|
|
height: base.height.default,
|
|
|
|
},
|
|
|
|
});
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
|
|
|
const createImg2ImgSlice: StateCreator<OnnxState, [], [], Img2ImgSlice> = (set) => ({
|
|
|
|
img2img: {
|
|
|
|
...defaults,
|
2023-01-13 20:39:07 +00:00
|
|
|
source: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
strength: base.strength.default,
|
|
|
|
},
|
|
|
|
setImg2Img(params) {
|
|
|
|
set((prev) => ({
|
|
|
|
img2img: {
|
|
|
|
...prev.img2img,
|
|
|
|
...params,
|
|
|
|
},
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
resetImg2Img() {
|
|
|
|
set({
|
|
|
|
img2img: {
|
|
|
|
...defaults,
|
2023-01-13 20:39:07 +00:00
|
|
|
source: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
strength: base.strength.default,
|
|
|
|
},
|
|
|
|
});
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
|
|
|
const createInpaintSlice: StateCreator<OnnxState, [], [], InpaintSlice> = (set) => ({
|
|
|
|
inpaint: {
|
|
|
|
...defaults,
|
2023-01-15 17:16:39 +00:00
|
|
|
filter: 'none',
|
2023-01-13 20:39:07 +00:00
|
|
|
mask: null,
|
2023-01-15 17:16:39 +00:00
|
|
|
noise: 'histogram',
|
2023-01-15 17:10:02 +00:00
|
|
|
source: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
},
|
|
|
|
setInpaint(params) {
|
|
|
|
set((prev) => ({
|
|
|
|
inpaint: {
|
|
|
|
...prev.inpaint,
|
|
|
|
...params,
|
|
|
|
},
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
resetInpaint() {
|
|
|
|
set({
|
|
|
|
inpaint: {
|
|
|
|
...defaults,
|
2023-01-15 17:16:39 +00:00
|
|
|
filter: 'none',
|
2023-01-13 20:39:07 +00:00
|
|
|
mask: null,
|
2023-01-15 17:16:39 +00:00
|
|
|
noise: 'histogram',
|
2023-01-13 20:39:07 +00:00
|
|
|
source: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
},
|
|
|
|
});
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
|
|
|
const createHistorySlice: StateCreator<OnnxState, [], [], HistorySlice> = (set) => ({
|
|
|
|
history: [],
|
|
|
|
limit: 4,
|
2023-01-13 03:12:20 +00:00
|
|
|
loading: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
pushHistory(image) {
|
|
|
|
set((prev) => ({
|
|
|
|
...prev,
|
|
|
|
history: [
|
|
|
|
image,
|
|
|
|
...prev.history,
|
|
|
|
],
|
2023-01-13 03:12:20 +00:00
|
|
|
loading: null,
|
2023-01-12 23:49:51 +00:00
|
|
|
}));
|
|
|
|
},
|
|
|
|
removeHistory(image) {
|
2023-01-12 23:51:34 +00:00
|
|
|
set((prev) => ({
|
|
|
|
...prev,
|
|
|
|
history: prev.history.filter((it) => it.output !== image.output),
|
|
|
|
}));
|
2023-01-12 23:49:51 +00:00
|
|
|
},
|
|
|
|
setLimit(limit) {
|
|
|
|
set((prev) => ({
|
|
|
|
...prev,
|
|
|
|
limit,
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
setLoading(loading) {
|
|
|
|
set((prev) => ({
|
|
|
|
...prev,
|
|
|
|
loading,
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
2023-01-15 00:07:05 +00:00
|
|
|
const createOutpaintSlice: StateCreator<OnnxState, [], [], OutpaintSlice> = (set) => ({
|
|
|
|
outpaint: {
|
2023-01-15 00:30:27 +00:00
|
|
|
enabled: false,
|
2023-01-15 00:07:05 +00:00
|
|
|
left: 0,
|
|
|
|
right: 0,
|
|
|
|
top: 0,
|
|
|
|
bottom: 0,
|
|
|
|
},
|
|
|
|
setOutpaint(pixels) {
|
|
|
|
set((prev) => ({
|
|
|
|
outpaint: {
|
|
|
|
...prev.outpaint,
|
|
|
|
...pixels,
|
|
|
|
}
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
|
|
|
const createBrushSlice: StateCreator<OnnxState, [], [], BrushSlice> = (set) => ({
|
|
|
|
brush: {
|
|
|
|
color: 255,
|
|
|
|
size: 8,
|
|
|
|
strength: 0.5,
|
|
|
|
},
|
|
|
|
setBrush(brush) {
|
|
|
|
set((prev) => ({
|
|
|
|
brush: {
|
|
|
|
...prev.brush,
|
|
|
|
...brush,
|
|
|
|
},
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
2023-01-16 20:05:04 +00:00
|
|
|
const createUpscaleSlice: StateCreator<OnnxState, [], [], UpscaleSlice> = (set) => ({
|
|
|
|
upscale: {
|
|
|
|
denoise: 0.5,
|
|
|
|
enabled: false,
|
|
|
|
faces: false,
|
|
|
|
scale: 1,
|
2023-01-16 20:52:56 +00:00
|
|
|
outscale: 1,
|
2023-01-17 04:54:01 +00:00
|
|
|
faceStrength: 0.5,
|
2023-01-16 20:05:04 +00:00
|
|
|
},
|
|
|
|
setUpscale(upscale) {
|
|
|
|
set((prev) => ({
|
|
|
|
upscale: {
|
|
|
|
...prev.upscale,
|
|
|
|
...upscale,
|
|
|
|
}
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
const createDefaultSlice: StateCreator<OnnxState, [], [], DefaultSlice> = (set) => ({
|
|
|
|
defaults: {
|
|
|
|
...defaults,
|
|
|
|
},
|
|
|
|
setDefaults(params) {
|
|
|
|
set((prev) => ({
|
|
|
|
defaults: {
|
|
|
|
...prev.defaults,
|
|
|
|
...params,
|
|
|
|
}
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
2023-01-17 02:11:10 +00:00
|
|
|
const createModelSlice: StateCreator<OnnxState, [], [], ModelSlice> = (set) => ({
|
|
|
|
model: {
|
|
|
|
model: '',
|
|
|
|
platform: '',
|
|
|
|
upscaling: '',
|
|
|
|
correction: '',
|
|
|
|
},
|
|
|
|
setModel(params) {
|
|
|
|
set((prev) => ({
|
|
|
|
model: {
|
|
|
|
...prev.model,
|
|
|
|
...params,
|
|
|
|
}
|
|
|
|
}));
|
|
|
|
},
|
|
|
|
});
|
|
|
|
|
2023-01-12 23:49:51 +00:00
|
|
|
return {
|
2023-01-16 20:05:04 +00:00
|
|
|
createBrushSlice,
|
2023-01-12 23:49:51 +00:00
|
|
|
createDefaultSlice,
|
|
|
|
createHistorySlice,
|
|
|
|
createImg2ImgSlice,
|
|
|
|
createInpaintSlice,
|
2023-01-17 02:11:10 +00:00
|
|
|
createModelSlice,
|
2023-01-15 00:07:05 +00:00
|
|
|
createOutpaintSlice,
|
2023-01-16 20:05:04 +00:00
|
|
|
createTxt2ImgSlice,
|
|
|
|
createUpscaleSlice,
|
2023-01-12 23:49:51 +00:00
|
|
|
};
|
|
|
|
}
|
2023-01-13 00:10:46 +00:00
|
|
|
|
|
|
|
export const ClientContext = createContext<Maybe<ApiClient>>(undefined);
|
2023-01-17 02:11:10 +00:00
|
|
|
export const ConfigContext = createContext<Maybe<ConfigParams>>(undefined);
|
2023-01-13 00:10:46 +00:00
|
|
|
export const StateContext = createContext<Maybe<StoreApi<OnnxState>>>(undefined);
|