From 1dae8eea6f13aa785968e237d68cc85f0be6d00d Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Mon, 8 Jan 2024 20:18:16 -0600
Subject: [PATCH] feat(gui): filter pipelines based on tab and set a reasonable
 default

---
 gui/src/components/card/LoadingCard.tsx     | 11 +++--
 gui/src/components/control/ModelControl.tsx | 48 ++++++++++++++++++++-
 gui/src/components/tab/Img2Img.tsx          |  3 +-
 gui/src/components/tab/Inpaint.tsx          |  3 +-
 gui/src/components/tab/Txt2Img.tsx          |  3 +-
 gui/src/components/tab/Upscale.tsx          |  3 +-
 6 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/gui/src/components/card/LoadingCard.tsx b/gui/src/components/card/LoadingCard.tsx
index 18921dbc..e3383d0c 100644
--- a/gui/src/components/card/LoadingCard.tsx
+++ b/gui/src/components/card/LoadingCard.tsx
@@ -11,6 +11,7 @@ import { shallow } from 'zustand/shallow';
 import { POLL_TIME } from '../../config.js';
 import { ClientContext, ConfigContext, OnnxState, StateContext } from '../../state/full.js';
 import { JobResponse, JobStatus } from '../../types/api-v2.js';
+import { visibleIndex } from '../../utils.js';
 
 const LOADING_PERCENT = 100;
 const LOADING_OVERAGE = 99;
@@ -147,11 +148,15 @@ function getStatus(data: Maybe<Array<JobResponse>>) {
 
 function getQueue(data: Maybe<Array<JobResponse>>) {
   if (doesExist(data) && data[0].status === JobStatus.PENDING) {
-    return data[0].queue;
+    const { current, total } = data[0].queue;
+    return {
+      current: visibleIndex(current),
+      total: total.toFixed(0),
+    };
   }
 
   return {
-    current: 0,
-    total: 0,
+    current: '0',
+    total: '0',
   };
 }
diff --git a/gui/src/components/control/ModelControl.tsx b/gui/src/components/control/ModelControl.tsx
index 54b99846..20881b9a 100644
--- a/gui/src/components/control/ModelControl.tsx
+++ b/gui/src/components/control/ModelControl.tsx
@@ -9,15 +9,17 @@ import { STALE_TIME } from '../../config.js';
 import { ClientContext } from '../../state/full.js';
 import { ModelParams } from '../../types/params.js';
 import { QueryList } from '../input/QueryList.js';
+import { JobType } from '../../types/api-v2.js';
 
 export interface ModelControlProps {
   model: ModelParams;
   setModel(params: Partial<ModelParams>): void;
+  tab: JobType;
 }
 
 export function ModelControl(props: ModelControlProps) {
   // eslint-disable-next-line @typescript-eslint/unbound-method
-  const { model, setModel } = props;
+  const { model, setModel, tab } = props;
 
   const client = mustExist(useContext(ClientContext));
   const { t } = useTranslation();
@@ -54,6 +56,7 @@ export function ModelControl(props: ModelControlProps) {
       name={t('parameter.pipeline')}
       query={{
         result: pipelines,
+        selector: (result) => filterValidPipelines(result, tab),
       }}
       value={model.pipeline}
       onChange={(pipeline) => {
@@ -113,3 +116,46 @@ export function ModelControl(props: ModelControlProps) {
     >{t('admin.restart')}</Button>
   </Stack>;
 }
+
+// plugin pipelines will show up on all tabs for now
+export const PIPELINE_TABS: Record<string, Array<JobType>> = {
+  'txt2img': [JobType.TXT2IMG],
+  'txt2img-sdxl': [JobType.TXT2IMG],
+  'panorama': [JobType.TXT2IMG, JobType.IMG2IMG],
+  'panorama-sdxl': [JobType.TXT2IMG, JobType.IMG2IMG],
+  'lpw': [JobType.TXT2IMG, JobType.IMG2IMG, JobType.INPAINT],
+  'img2img': [JobType.IMG2IMG],
+  'img2img-sdxl': [JobType.IMG2IMG],
+  'controlnet': [JobType.IMG2IMG],
+  'pix2pix': [JobType.IMG2IMG],
+  'inpaint': [JobType.INPAINT],
+  'upscale': [JobType.UPSCALE],
+};
+
+export const DEFAULT_PIPELINE: Record<JobType, string> = {
+  [JobType.TXT2IMG]: 'txt2img',
+  [JobType.IMG2IMG]: 'img2img',
+  [JobType.INPAINT]: 'inpaint',
+  [JobType.UPSCALE]: 'upscale',
+  [JobType.BLEND]: '',
+  [JobType.CHAIN]: '',
+};
+
+export const FIRST_A = -1;
+export const FIRST_B = +1;
+
+export function filterValidPipelines(pipelines: Array<string>, tab: JobType): Array<string> {
+  const defaultPipeline = DEFAULT_PIPELINE[tab];
+  return pipelines.filter((pipeline) => PIPELINE_TABS[pipeline].includes(tab)).sort((a, b) => {
+    // put validPipelines.default first
+    if (a === defaultPipeline) {
+      return FIRST_A;
+    }
+
+    if (b === defaultPipeline) {
+      return FIRST_B;
+    }
+
+    return a.localeCompare(b);
+  });
+}
diff --git a/gui/src/components/tab/Img2Img.tsx b/gui/src/components/tab/Img2Img.tsx
index 61911f4b..d764d14d 100644
--- a/gui/src/components/tab/Img2Img.tsx
+++ b/gui/src/components/tab/Img2Img.tsx
@@ -19,6 +19,7 @@ import { UpscaleControl } from '../control/UpscaleControl.js';
 import { ImageInput } from '../input/ImageInput.js';
 import { NumericField } from '../input/NumericField.js';
 import { QueryList } from '../input/QueryList.js';
+import { JobType } from '../../types/api-v2.js';
 
 export function Img2Img() {
   const { params } = mustExist(useContext(ConfigContext));
@@ -67,7 +68,7 @@ export function Img2Img() {
         setParams={setImg2Img}
         setUpscale={setUpscale}
       />
-      <ModelControl model={model} setModel={setModel} />
+      <ModelControl model={model} setModel={setModel} tab={JobType.IMG2IMG} />
       <ImageInput
         filter={IMAGE_FILTER}
         image={source}
diff --git a/gui/src/components/tab/Inpaint.tsx b/gui/src/components/tab/Inpaint.tsx
index 87ea7405..601e0af7 100644
--- a/gui/src/components/tab/Inpaint.tsx
+++ b/gui/src/components/tab/Inpaint.tsx
@@ -21,6 +21,7 @@ import { ImageInput } from '../input/ImageInput.js';
 import { MaskCanvas } from '../input/MaskCanvas.js';
 import { NumericField } from '../input/NumericField.js';
 import { QueryList } from '../input/QueryList.js';
+import { JobType } from '../../types/api-v2.js';
 
 export function Inpaint() {
   const { params } = mustExist(useContext(ConfigContext));
@@ -99,7 +100,7 @@ export function Inpaint() {
         setParams={setInpaint}
         setUpscale={setUpscale}
       />
-      <ModelControl model={model} setModel={setModel} />
+      <ModelControl model={model} setModel={setModel} tab={JobType.INPAINT} />
       {renderBanner()}
       <ImageInput
         filter={IMAGE_FILTER}
diff --git a/gui/src/components/tab/Txt2Img.tsx b/gui/src/components/tab/Txt2Img.tsx
index 060e6766..e344e3c0 100644
--- a/gui/src/components/tab/Txt2Img.tsx
+++ b/gui/src/components/tab/Txt2Img.tsx
@@ -18,6 +18,7 @@ import { ModelControl } from '../control/ModelControl.js';
 import { UpscaleControl } from '../control/UpscaleControl.js';
 import { VariableControl } from '../control/VariableControl.js';
 import { NumericField } from '../input/NumericField.js';
+import { JobType } from '../../types/api-v2.js';
 
 export function SizeControl() {
   const { params } = mustExist(useContext(ConfigContext));
@@ -98,7 +99,7 @@ export function Txt2Img() {
         setParams={setParams}
         setUpscale={setUpscale}
       />
-      <ModelControl model={model} setModel={setModel} />
+      <ModelControl model={model} setModel={setModel} tab={JobType.TXT2IMG} />
       <ImageControl selector={selectParams} onChange={setParams} />
       <SizeControl />
       <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
diff --git a/gui/src/components/tab/Upscale.tsx b/gui/src/components/tab/Upscale.tsx
index 4c0ea2e1..aa9680da 100644
--- a/gui/src/components/tab/Upscale.tsx
+++ b/gui/src/components/tab/Upscale.tsx
@@ -17,6 +17,7 @@ import { ModelControl } from '../control/ModelControl.js';
 import { UpscaleControl } from '../control/UpscaleControl.js';
 import { ImageInput } from '../input/ImageInput.js';
 import { PromptInput } from '../input/PromptInput.js';
+import { JobType } from '../../types/api-v2.js';
 
 export function Upscale() {
   async function uploadSource() {
@@ -53,7 +54,7 @@ export function Upscale() {
         setParams={setParams}
         setUpscale={setUpscale}
       />
-      <ModelControl model={model} setModel={setModel} />
+      <ModelControl model={model} setModel={setModel} tab={JobType.UPSCALE} />
       <ImageInput
         filter={IMAGE_FILTER}
         image={params.source}