From 46e047bd0a9cc784ae94263f3f1db8be6695f07c Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Thu, 5 Jan 2023 22:33:06 -0600
Subject: [PATCH] feat(gui): load models from server

---
 gui/src/api/client.ts          | 21 ++++++++++++++++++++-
 gui/src/components/OnnxWeb.tsx | 30 ++++++++++++++++++++++++------
 gui/src/components/Txt2Img.tsx | 10 ++++++++--
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/gui/src/api/client.ts b/gui/src/api/client.ts
index 2882ff0e..46a87c52 100644
--- a/gui/src/api/client.ts
+++ b/gui/src/api/client.ts
@@ -1,13 +1,17 @@
 import { doesExist } from '@apextoaster/js-utils';
 
 export interface Txt2ImgParams {
+  model?: string;
+  platform?: string;
+  scheduler?: string;
+
   prompt: string;
   cfg: number;
   steps: number;
+
   width?: number;
   height?: number;
   seed?: string;
-  scheduler?: string;
 }
 
 export interface ApiResponse {
@@ -16,8 +20,10 @@ export interface ApiResponse {
 }
 
 export interface ApiClient {
+  models(): Promise<Array<string>>;
   platforms(): Promise<Array<string>>;
   schedulers(): Promise<Array<string>>;
+
   txt2img(params: Txt2ImgParams): Promise<ApiResponse>;
 }
 
@@ -40,6 +46,11 @@ export function makeClient(root: string, f = fetch): ApiClient {
   let pending: Promise<ApiResponse> | undefined;
 
   return {
+    async models(): Promise<Array<string>> {
+      const path = new URL('/settings/models', root);
+      const res = await f(path);
+      return await res.json() as Array<string>;
+    },
     async schedulers(): Promise<Array<string>> {
       const path = new URL('/settings/schedulers', root);
       const res = await f(path);
@@ -71,6 +82,14 @@ export function makeClient(root: string, f = fetch): ApiClient {
         url.searchParams.append('seed', params.seed);
       }
 
+      if (doesExist(params.model)) {
+        url.searchParams.append('model', params.model);
+      }
+
+      if (doesExist(params.platform)) {
+        url.searchParams.append('platform', params.platform);
+      }
+
       if (doesExist(params.scheduler)) {
         url.searchParams.append('scheduler', params.scheduler);
       }
diff --git a/gui/src/components/OnnxWeb.tsx b/gui/src/components/OnnxWeb.tsx
index 47423ccd..67d2307e 100644
--- a/gui/src/components/OnnxWeb.tsx
+++ b/gui/src/components/OnnxWeb.tsx
@@ -1,9 +1,11 @@
+import { mustExist } from '@apextoaster/js-utils';
 import { TabContext, TabList, TabPanel } from '@mui/lab';
 import { Box, Container, MenuItem, Select, Tab, Typography } from '@mui/material';
 import * as React from 'react';
+import { useQuery } from 'react-query';
 
 import { ApiClient } from '../api/client.js';
-import { Txt2Img } from './Txt2Img.js';
+import { STALE_TIME, Txt2Img } from './Txt2Img.js';
 
 const { useState } = React;
 
@@ -13,22 +15,38 @@ export interface OnnxWebProps {
 
 export function OnnxWeb(props: OnnxWebProps) {
   const [tab, setTab] = useState('1');
-  const [model, setModel] = useState('v1.5');
+  const [model, setModel] = useState('stable-diffusion-onnx-v1-5');
+
+  const models = useQuery('models', async () => props.client.models(), {
+    staleTime: STALE_TIME,
+  });
+
+  function renderModels() {
+    switch (models.status) {
+      case 'error':
+        return <MenuItem value='error'>Error</MenuItem>;
+      case 'loading':
+        return <MenuItem value='loading'>Loading</MenuItem>;
+      case 'success':
+        return mustExist(models.data).map((name) => <MenuItem key={name} value={name}>{name}</MenuItem>);
+      default:
+        return <MenuItem value='error'>Unknown Error</MenuItem>;
+    }
+  }
 
   return (
     <div>
       <Container>
         <Box sx={{ my: 4 }}>
           <Typography variant='h3' gutterBottom>
-            ONNX Web GUI
+            ONNX Web
           </Typography>
         </Box>
         <Box sx={{ my: 4 }}>
           <Select value={model} onChange={(e) => {
             setModel(e.target.value);
           }}>
-            <MenuItem value='v1.4'>Stable Diffusion v1.4</MenuItem>
-            <MenuItem value='v1.5'>Stable Diffusion v1.5</MenuItem>
+            {renderModels()}
           </Select>
         </Box>
         <TabContext value={tab}>
@@ -42,7 +60,7 @@ export function OnnxWeb(props: OnnxWebProps) {
             </TabList>
           </Box>
           <TabPanel value="1">
-            <Txt2Img {...props} />
+            <Txt2Img client={props.client} model={model} />
           </TabPanel>
           <TabPanel value="2">
             <Box>
diff --git a/gui/src/components/Txt2Img.tsx b/gui/src/components/Txt2Img.tsx
index 2ef543a0..75bd4996 100644
--- a/gui/src/components/Txt2Img.tsx
+++ b/gui/src/components/Txt2Img.tsx
@@ -8,7 +8,7 @@ import { ImageControl, ImageParams } from './ImageControl.js';
 
 const { useState } = React;
 
-const STALE_TIME = 3_000;
+export const STALE_TIME = 3_000;
 
 // TODO: set up i18next
 const PLATFORM_NAMES: Record<string, string> = {
@@ -28,13 +28,19 @@ const SCHEDULER_NAMES: Record<string, string> = {
 
 export interface Txt2ImgProps {
   client: ApiClient;
+  model: string;
 }
 
 export function Txt2Img(props: Txt2ImgProps) {
   const { client } = props;
 
   async function generateImage() {
-    return client.txt2img({ ...params, prompt, scheduler });
+    return client.txt2img({
+      ...params,
+      model: props.model,
+      prompt,
+      scheduler,
+    });
   }
 
   const generate = useMutation(generateImage);