add VAE selection

public/scripts/extensions/stable-diffusion/comfyWorkflowEditor.html

@@ -12,6 +12,7 @@
 					<li data-placeholder="prompt" class="sd_comfy_workflow_editor_not_found">"%prompt%"</li>
 					<li data-placeholder="negative_prompt" class="sd_comfy_workflow_editor_not_found">"%negative_prompt%"</li>
 					<li data-placeholder="model" class="sd_comfy_workflow_editor_not_found">"%model%"</li>
+					<li data-placeholder="vae" class="sd_comfy_workflow_editor_not_found">"%vae%"</li>
 					<li data-placeholder="sampler" class="sd_comfy_workflow_editor_not_found">"%sampler%"</li>
 					<li data-placeholder="scheduler" class="sd_comfy_workflow_editor_not_found">"%scheduler%"</li>
 					<li data-placeholder="steps" class="sd_comfy_workflow_editor_not_found">"%steps%"</li>

public/scripts/extensions/stable-diffusion/index.js

@@ -186,6 +186,7 @@ const defaultSettings = {
     negative_prompt: defaultNegative,
     sampler: 'DDIM',
     model: '',
+    vae: '',
 
     // Automatic1111/Horde exclusives
     restore_faces: false,
@@ -480,7 +481,7 @@ async function loadSettings() {
     toggleSourceControls();
     addPromptTemplates();
 
-    await Promise.all([loadSamplers(), loadModels(), loadSchedulers()]);
+    await Promise.all([loadSamplers(), loadModels(), loadSchedulers(), loadVaes()]);
 }
 
 function addPromptTemplates() {
@@ -914,6 +915,7 @@ async function validateComfyUrl() {
         await loadSamplers();
         await loadSchedulers();
         await loadModels();
+        await loadVaes();
         toastr.success('ComfyUI API connected.');
     } catch (error) {
         toastr.error(`Could not validate ComfyUI API: ${error.message}`);
@@ -960,6 +962,10 @@ async function getAutoRemoteModel() {
     }
 }
 
+async function onVaeChange() {
+    extension_settings.sd.vae = $('#sd_vae').find(':selected').val();
+}
+
 async function getAutoRemoteUpscalers() {
     try {
         const result = await fetch('/api/sd/upscalers', {
@@ -1473,6 +1479,65 @@ async function loadComfySchedulers() {
     }
 }
 
+async function loadVaes() {
+    $('#sd_vae').empty();
+    let vaes = [];
+
+    switch (extension_settings.sd.source) {
+        case sources.extras:
+            vaes = ['N/A'];
+            break;
+        case sources.horde:
+            vaes = ['N/A'];
+            break;
+        case sources.auto:
+            vaes = ['N/A'];
+            break;
+        case sources.novel:
+            vaes = ['N/A'];
+            break;
+        case sources.vlad:
+            vaes = ['N/A'];
+            break;
+        case sources.openai:
+            vaes = ['N/A'];
+            break;
+        case sources.comfy:
+            vaes = await loadComfyVaes();
+            break;
+    }
+
+    for (const vae of vaes) {
+        const option = document.createElement('option');
+        option.innerText = vae;
+        option.value = vae;
+        option.selected = vae === extension_settings.sd.vae;
+        $('#sd_vae').append(option);
+    }
+}
+
+async function loadComfyVaes() {
+    if (!extension_settings.sd.comfy_url) {
+        return [];
+    }
+
+    try {
+        const result = await fetch(`/api/sd/comfy/vaes`, {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_url,
+            })
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+        return await result.json();
+    } catch (error) {
+        return [];
+    }
+}
+
 function getGenerationType(prompt) {
     let mode = generationMode.FREE;
 
@@ -2037,6 +2102,7 @@ async function generateComfyImage(prompt) {
     const placeholders = [
         'negative_prompt',
         'model',
+        'vae',
         'sampler',
         'scheduler',
         'steps',
@@ -2294,6 +2360,7 @@ jQuery(async () => {
     $('#sd_scale').on('input', onScaleInput);
     $('#sd_steps').on('input', onStepsInput);
     $('#sd_model').on('change', onModelChange);
+    $('#sd_vae').on('change', onVaeChange);
     $('#sd_sampler').on('change', onSamplerChange);
     $('#sd_scheduler').on('change', onSchedulerChange);
     $('#sd_prompt_prefix').on('input', onPromptPrefixInput);

public/scripts/extensions/stable-diffusion/settings.html

@@ -150,6 +150,10 @@
                 <label for="sd_scheduler">Scheduler</label>
                 <select id="sd_scheduler"></select>
             </div>
+            <div data-sd-source="comfy">
+                <label for="sd_vae">VAE</label>
+                <select id="sd_vae"></select>
+            </div>
             <div class="flex-container marginTop10 margin-bot-10px">
                 <label class="flex1 checkbox_label">
                     <input id="sd_restore_faces" type="checkbox" />

src/stable-diffusion.js

@@ -418,6 +418,24 @@ function registerEndpoints(app, jsonParser) {
         }
     });
 
+    app.post('/api/sd/comfy/vaes', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/object_info'
+
+			const result = await fetch(url);
+			if (!result.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+
+			const data = await result.json();
+			return response.send(data.VAELoader.input.required.vae_name[0]);
+		} catch (error) {
+			console.log(error);
+			return response.sendStatus(500);
+		}
+	});
+
     app.post('/api/sd/comfy/generate', jsonParser, async (request, response) => {
         try {
             const url = new URL(request.body.url);