onnx-web/gui/src/main.tsx

/* eslint-disable no-console */
import { doesExist, mustExist } from '@apextoaster/js-utils';
import { merge } from 'lodash';
import * as React from 'react';
import ReactDOM from 'react-dom/client';
import { QueryClient, QueryClientProvider } from 'react-query';
import { createStore } from 'zustand';
import { createJSONStorage, persist } from 'zustand/middleware';

import { makeClient } from './client.js';
import { OnnxError } from './components/OnnxError.js';
import { OnnxWeb } from './components/OnnxWeb.js';
import { Config, loadConfig } from './config.js';
import { ClientContext, createStateSlices, OnnxState, StateContext } from './state.js';

export function getApiRoot(config: Config): string {
  const query = new URLSearchParams(window.location.search);
  const api = query.get('api');

  if (doesExist(api)) {
    return api;
  } else {
    return config.api.root;
  }
}

export async function main() {
  // load config from GUI server
  const config = await loadConfig();

  // use that to create an API client
  const root = getApiRoot(config);
  const client = makeClient(root);

  // prep react-dom
  const appElement = mustExist(document.getElementById('app'));
  const app = ReactDOM.createRoot(appElement);

  try {
    // load full params from the API server and merge with the initial client config
    const params = await client.params();
    merge(params, config.params);

    // prep zustand with a slice for each tab, using local storage
    const {
      createBrushSlice,
      createDefaultSlice,
      createHistorySlice,
      createImg2ImgSlice,
      createInpaintSlice,
      createOutpaintSlice,
      createTxt2ImgSlice,
      createUpscaleSlice,
    } = createStateSlices(params);
    const state = createStore<OnnxState, [['zustand/persist', OnnxState]]>(persist((...slice) => ({
      ...createBrushSlice(...slice),
      ...createDefaultSlice(...slice),
      ...createHistorySlice(...slice),
      ...createImg2ImgSlice(...slice),
      ...createInpaintSlice(...slice),
      ...createTxt2ImgSlice(...slice),
      ...createOutpaintSlice(...slice),
      ...createUpscaleSlice(...slice),
    }), {
      name: 'onnx-web',
      partialize(s) {
        return {
          ...s,
          img2img: {
            ...s.img2img,
            source: undefined,
          },
          inpaint: {
            ...s.inpaint,
            mask: undefined,
            source: undefined,
          },
        };
      },
      storage: createJSONStorage(() => localStorage),
      version: 3,
    }));

    // prep react-query client
    const query = new QueryClient();

    // go
    app.render(<QueryClientProvider client={query}>
      <ClientContext.Provider value={client}>
        <StateContext.Provider value={state}>
          <OnnxWeb client={client} config={params} />
        </StateContext.Provider>
      </ClientContext.Provider>
    </QueryClientProvider>);
  } catch (err) {
    app.render(<OnnxError root={root} />);
  }
}

window.addEventListener('load', () => {
  console.log('launching onnx-web');
  main().catch((err) => {
    console.error('error in main', err);
  });
}, false);