fix(gui): add state migrations for new unet/vae params (#427)

gui/src/components/card/ImageCard.tsx

@@ -8,9 +8,10 @@ import { useHash } from 'react-use/lib/useHash';
 import { useStore } from 'zustand';
 import { shallow } from 'zustand/shallow';
 
-import { BLEND_SOURCES, ConfigContext, OnnxState, StateContext } from '../../state/full.js';
+import { ConfigContext, OnnxState, StateContext } from '../../state/full.js';
 import { ImageResponse } from '../../types/api.js';
 import { range, visibleIndex } from '../../utils.js';
+import { BLEND_SOURCES } from '../../constants.js';
 
 export interface ImageCardProps {
   image: ImageResponse;

gui/src/components/tab/Blend.tsx

@@ -8,7 +8,8 @@ import { useStore } from 'zustand';
 import { shallow } from 'zustand/shallow';
 
 import { IMAGE_FILTER } from '../../config.js';
-import { BLEND_SOURCES, ClientContext, OnnxState, StateContext } from '../../state/full.js';
+import { BLEND_SOURCES } from '../../constants.js';
+import { ClientContext, OnnxState, StateContext } from '../../state/full.js';
 import { TabState } from '../../state/types.js';
 import { BlendParams, BrushParams, ModelParams, UpscaleParams } from '../../types/params.js';
 import { range } from '../../utils.js';

gui/src/constants.ts

@@ -0,0 +1,31 @@
+
+export const BLEND_SOURCES = 2;
+
+/**
+ * Default parameters for the inpaint brush.
+ *
+ * Not provided by the server yet.
+ */
+export const DEFAULT_BRUSH = {
+  color: 255,
+  size: 8,
+  strength: 0.5,
+};
+
+/**
+ * Default parameters for the image history.
+ *
+ * Not provided by the server yet.
+ */
+export const DEFAULT_HISTORY = {
+  /**
+   * The number of images to be shown.
+   */
+  limit: 4,
+
+  /**
+   * The number of additional images to be kept in history, so they can scroll
+   * back into view when you delete one. Does not include deleted images.
+   */
+  scrollback: 2,
+};

gui/src/main.tsx

@@ -30,6 +30,7 @@ import {
   StateContext,
 } from './state/full.js';
 import { I18N_STRINGS } from './strings/all.js';
+import { applyStateMigrations, UnknownState } from './state/migration/default.js';
 
 export const INITIAL_LOAD_TIMEOUT = 5_000;
 
@@ -70,6 +71,9 @@ export async function renderApp(config: Config, params: ServerParams, logger: Lo
     ...createResetSlice(...slice),
     ...createProfileSlice(...slice),
   }), {
+    migrate(persistedState, version) {
+      return applyStateMigrations(params, persistedState as UnknownState, version);
+    },
     name: STATE_KEY,
     partialize(s) {
       return {

gui/src/state/blend.ts

@@ -1,10 +1,11 @@
+import { DEFAULT_BRUSH } from '../constants.js';
 import {
   BlendParams,
   BrushParams,
   ModelParams,
   UpscaleParams,
 } from '../types/params.js';
-import { DEFAULT_BRUSH, Slice, TabState } from './types.js';
+import { Slice, TabState } from './types.js';
 
 export interface BlendSlice {
   blend: TabState<BlendParams>;

gui/src/state/full.ts

@@ -69,9 +69,7 @@ export const STATE_KEY = 'onnx-web';
 /**
  * Current state version for zustand persistence.
  */
-export const STATE_VERSION = 7;
-
-export const BLEND_SOURCES = 2;
+export const STATE_VERSION = 11;
 
 export function baseParamsFromServer(defaults: ServerParams): Required<BaseImgParams> {
   return {

gui/src/state/history.ts

@@ -1,6 +1,7 @@
 import { Maybe } from '@apextoaster/js-utils';
 import { ImageResponse, ReadyResponse, RetryParams } from '../types/api.js';
-import { DEFAULT_HISTORY, Slice } from './types.js';
+import { Slice } from './types.js';
+import { DEFAULT_HISTORY } from '../constants.js';
 
 export interface HistoryItem {
   image: ImageResponse;

gui/src/state/inpaint.ts

@@ -1,4 +1,5 @@
 import { ServerParams } from '../config.js';
+import { DEFAULT_BRUSH } from '../constants.js';
 import {
   BaseImgParams,
   BrushParams,
@@ -8,7 +9,7 @@ import {
   OutpaintPixels,
   UpscaleParams,
 } from '../types/params.js';
-import { DEFAULT_BRUSH, Slice, TabState } from './types.js';
+import { Slice, TabState } from './types.js';
 export interface InpaintSlice {
   inpaint: TabState<InpaintParams>;
   inpaintBrush: BrushParams;

gui/src/state/migration/default.ts

@@ -0,0 +1,82 @@
+/* eslint-disable camelcase */
+import { ServerParams } from '../../config.js';
+import { BaseImgParams } from '../../types/params.js';
+import { OnnxState, STATE_VERSION } from '../full.js';
+import { Img2ImgSlice } from '../img2img.js';
+import { InpaintSlice } from '../inpaint.js';
+import { Txt2ImgSlice } from '../txt2img.js';
+import { UpscaleSlice } from '../upscale.js';
+
+export const REMOVE_KEYS = ['tile', 'overlap'] as const;
+
+export type RemovedKeys = typeof REMOVE_KEYS[number];
+
+// TODO: can the compiler calculate this?
+export type AddedKeysV11 = 'unet_tile' | 'unet_overlap' | 'vae_tile' | 'vae_overlap';
+
+export type BaseImgParamsV7<T extends BaseImgParams> = Omit<T, AddedKeysV11> & {
+  overlap: number;
+  tile: number;
+};
+
+export type OnnxStateV7 = Omit<OnnxState, 'img2img' | 'txt2img'> & {
+  img2img: BaseImgParamsV7<Img2ImgSlice['img2img']>;
+  inpaint: BaseImgParamsV7<InpaintSlice['inpaint']>;
+  txt2img: BaseImgParamsV7<Txt2ImgSlice['txt2img']>;
+  upscale: BaseImgParamsV7<UpscaleSlice['upscale']>;
+};
+
+export type PreviousState = OnnxStateV7;
+export type CurrentState = OnnxState;
+export type UnknownState = PreviousState | CurrentState;
+
+export function applyStateMigrations(params: ServerParams, previousState: UnknownState, version: number): OnnxState {
+  // eslint-disable-next-line no-console
+  console.log('applying migrations from %s to %s', version, STATE_VERSION);
+
+  if (version < STATE_VERSION) {
+    return migrateDefaults(params, previousState as PreviousState);
+  }
+
+  return previousState as CurrentState;
+}
+
+export function migrateDefaults(params: ServerParams, previousState: PreviousState): CurrentState {
+  // add any missing keys
+  const result: CurrentState = {
+    ...params,
+    ...previousState,
+    img2img: {
+      ...previousState.img2img,
+      unet_overlap: params.unet_overlap.default,
+      unet_tile: params.unet_tile.default,
+      vae_overlap: params.vae_overlap.default,
+      vae_tile: params.vae_tile.default,
+    },
+    inpaint: {
+      ...previousState.inpaint,
+      unet_overlap: params.unet_overlap.default,
+      unet_tile: params.unet_tile.default,
+      vae_overlap: params.vae_overlap.default,
+      vae_tile: params.vae_tile.default,
+    },
+    txt2img: {
+      ...previousState.txt2img,
+      unet_overlap: params.unet_overlap.default,
+      unet_tile: params.unet_tile.default,
+      vae_overlap: params.vae_overlap.default,
+      vae_tile: params.vae_tile.default,
+    },
+    upscale: {
+      ...previousState.upscale,
+      unet_overlap: params.unet_overlap.default,
+      unet_tile: params.unet_tile.default,
+      vae_overlap: params.vae_overlap.default,
+      vae_tile: params.vae_tile.default,
+    },
+  };
+
+  // TODO: remove extra keys
+
+  return result;
+}

gui/src/state/types.ts

@@ -15,32 +15,3 @@ export type TabState<TabParams> = ConfigFiles<Required<TabParams>> & ConfigState
  * Shorthand for state creator to reduce repeated arguments.
  */
 export type Slice<TState, TValue> = StateCreator<TState, [], [], TValue>;
-
-/**
- * Default parameters for the inpaint brush.
- *
- * Not provided by the server yet.
- */
-export const DEFAULT_BRUSH = {
-  color: 255,
-  size: 8,
-  strength: 0.5,
-};
-
-/**
- * Default parameters for the image history.
- *
- * Not provided by the server yet.
- */
-export const DEFAULT_HISTORY = {
-  /**
-   * The number of images to be shown.
-   */
-  limit: 4,
-
-  /**
-   * The number of additional images to be kept in history, so they can scroll
-   * back into view when you delete one. Does not include deleted images.
-   */
-  scrollback: 2,
-};