refresh the vae menu on when refreshing the UI

AbdullahAlfaraj · Jun 5, 2023 · 4a0d45b · 4a0d45b
1 parent d24a8e9
commit 4a0d45b
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 25 deletions.
diff --git a/after_detailer/src/after_detailer.tsx b/after_detailer/src/after_detailer.tsx
@@ -9,7 +9,6 @@ import {
     SpMenu,
     SpSliderWithLabel,
 } from '../../ultimate_sd_upscaler/src/elements'
-// import * as sdapi from '../../sdapi_py_re'
 
 import { AStore } from '../../main/src/astore'
 import { ui_config, model_list } from './config'
@@ -141,7 +140,6 @@ export class AfterDetailerComponent extends React.Component<{
                     </sp-label>
                     <button
                         className="btnSquare refreshButton"
-                        id="btnResetSettings"
                         title="Refresh the After Detailer Extension"
                         onClick={this.handleRefresh}
                     ></button>

diff --git a/index.js b/index.js
@@ -487,6 +487,7 @@ async function refreshUI() {
 
         g_controlnet_max_models = await control_net.requestControlNetMaxUnits()
         await control_net.initializeControlNetTab(g_controlnet_max_models)
+        await main.populateVAE()
     } catch (e) {
         console.warn(e)
     }

diff --git a/main/src/main.tsx b/main/src/main.tsx
@@ -6,7 +6,7 @@ import { SpMenu } from '../../ultimate_sd_upscaler/src/elements'
 
 import { getExtensionUrl } from '../../utility/sdapi/python_replacement'
 import * as api from '../../utility/api'
-
+declare let g_sd_url: string
 // class SDStore extends AStore {
 //     constructor(data: any) {
 //         super(data)
@@ -24,9 +24,10 @@ import * as api from '../../utility/api'
 
 const default_values: any = {
     vae_model_list: [],
+    current_vae: '',
 }
 
-export const sd_store = new AStore(default_values)
+export const store = new AStore(default_values)
 
 @observer
 export class VAEComponent extends React.Component<{
@@ -41,41 +42,63 @@ export class VAEComponent extends React.Component<{
             console.warn('changeVAEModel: vae_model: ', vae_model, e)
         }
     }
+    handleRefresh() {
+        populateVAE()
+    }
     render(): React.ReactNode {
         return (
-            <SpMenu
-                title="vae models"
-                items={sd_store.data.vae_model_list}
-                // disabled={script_store.disabled}
-                // style="width: 199px; margin-right: 5px"
-                label_item="Select A VAE"
-                // id={'model_list'}
-                onChange={(id: any, value: any) => {
-                    // script_store.setSelectedScript(value.item)
-                    console.log('onChange value: ', value)
-                    this.changeVAEModel(value.item)
-                }}
-            ></SpMenu>
+            <div style={{ display: 'flex' }}>
+                <SpMenu
+                    title="vae models"
+                    items={store.data.vae_model_list}
+                    // disabled={script_store.disabled}
+                    // style="width: 199px; margin-right: 5px"
+                    label_item="Select A VAE"
+                    // id={'model_list'}
+                    selected_index={store.data.vae_model_list.indexOf(
+                        store.data.current_vae
+                    )}
+                    onChange={(id: any, value: any) => {
+                        // script_store.setSelectedScript(value.item)
+                        console.log('onChange value: ', value)
+                        this.changeVAEModel(value.item)
+                    }}
+                ></SpMenu>
+                <button
+                    className="btnSquare refreshButton"
+                    title="Refresh VAE Models List"
+                    onClick={this.handleRefresh}
+                ></button>
+            </div>
         )
     }
 }
 const vaeContainerNode = document.getElementById('settingsVAEContainer')!
 const vaeRoot = ReactDOM.createRoot(vaeContainerNode)
 
-let vae_model_list = ['None']
-declare let g_sd_url: string
+// let vae_model_list = ['None']
+
+async function requestGetVAE() {
+    const full_url = `${g_sd_url}/sdapi/v1/options`
 
-async function populateVAE() {
+    const options = await api.requestGet(full_url)
+    return options?.sd_vae
+}
+export async function populateVAE() {
     const extension_url = getExtensionUrl()
 
     const full_url = `${extension_url}/vae/list`
-
     const vae_models = await api.requestGet(full_url)
+
     console.log('populateVAE vae_models: ', vae_models)
-    // vae_model_list = vae_models
-    sd_store.updateProperty('vae_model_list', vae_models)
+    store.updateProperty('vae_model_list', vae_models)
+
+    const current_vae = await requestGetVAE()
+    if (current_vae && vae_models.includes(current_vae)) {
+        store.updateProperty('current_vae', current_vae)
+    }
 }
-populateVAE()
+// populateVAE()
 
 vaeRoot.render(
     <React.StrictMode>

diff --git a/ultimate_sd_upscaler/src/ultimate_sd_upscaler.tsx b/ultimate_sd_upscaler/src/ultimate_sd_upscaler.tsx
@@ -178,7 +178,6 @@ export class UltimateSDUpscalerForm extends React.Component<{
                     </sp-label>
                     <button
                         className="btnSquare refreshButton"
-                        id="btnResetSettings"
                         title="Refresh the Ultimte SD Upscale script"
                         onClick={this.handleRefresh}
                     ></button>