From de6c8c1501e48459f39d796c07f431a62761aaee Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Tue, 31 Dec 2024 00:27:12 +0200
Subject: [PATCH] Remove forge override from non-forge SD requests

---
 src/endpoints/stable-diffusion.js | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/endpoints/stable-diffusion.js b/src/endpoints/stable-diffusion.js
index 694cbf843..4689fabe9 100644
--- a/src/endpoints/stable-diffusion.js
+++ b/src/endpoints/stable-diffusion.js
@@ -7,6 +7,7 @@ import sanitize from 'sanitize-filename';
 import { sync as writeFileAtomicSync } from 'write-file-atomic';
 import FormData from 'form-data';
 import urlJoin from 'url-join';
+import _ from 'lodash';
 
 import { delay, getBasicAuthHeader, tryParse } from '../util.js';
 import { jsonParser } from '../express-common.js';
@@ -293,23 +294,32 @@ router.post('/set-model', jsonParser, async (request, response) => {
 
 router.post('/generate', jsonParser, async (request, response) => {
     try {
-        console.log('SD WebUI request:', request.body);
+        try {
+            const optionsUrl = urlJoin(request.body.url, '/sdapi/v1/options');
+            const optionsResult = await fetch(optionsUrl, { headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } });
+            const optionsData = /** @type {any} */ (await optionsResult.json());
+            const isForge = 'forge_preset' in optionsData;
 
-        const url = new URL(request.body.url);
-        url.pathname = '/sdapi/v1/txt2img';
+            if (!isForge) {
+                _.unset(request.body, 'override_settings.forge_additional_modules');
+            }
+        } catch (error) {
+            console.log('SD WebUI failed to get options:', error);
+        }
 
         const controller = new AbortController();
         request.socket.removeAllListeners('close');
         request.socket.on('close', function () {
             if (!response.writableEnded) {
-                const url = new URL(request.body.url);
-                url.pathname = '/sdapi/v1/interrupt';
-                fetch(url, { method: 'POST', headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } });
+                const interruptUrl = urlJoin(request.body.url, '/sdapi/v1/interrupt');
+                fetch(interruptUrl, { method: 'POST', headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } });
             }
             controller.abort();
         });
 
-        const result = await fetch(url, {
+        console.log('SD WebUI request:', request.body);
+        const txt2imgUrl = urlJoin(request.body.url, '/sdapi/v1/txt2img');
+        const result = await fetch(txt2imgUrl, {
             method: 'POST',
             body: JSON.stringify(request.body),
             headers: {