From efcb56f0dc71caf2333dad4f13135e0f0390061e Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 21 Aug 2024 17:27:39 +0200
Subject: [PATCH] refac

---
 backend/apps/images/main.py                   | 25 +++++++++++++++
 src/lib/apis/images/index.ts                  | 32 +++++++++++++++++++
 .../components/admin/Settings/Images.svelte   | 27 +++++++++++++---
 3 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py
index 5244923d1..371b1c3bc 100644
--- a/backend/apps/images/main.py
+++ b/backend/apps/images/main.py
@@ -175,6 +175,31 @@ def get_automatic1111_api_auth():
         return f"Basic {auth1111_base64_encoded_string}"
 
 
+@app.get("/config/url/verify")
+async def verify_url(user=Depends(get_admin_user)):
+    if app.state.config.ENGINE == "automatic1111":
+        try:
+            r = requests.get(
+                url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
+                headers={"authorization": get_automatic1111_api_auth()},
+            )
+            r.raise_for_status()
+            return True
+        except Exception as e:
+            app.state.config.ENABLED = False
+            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
+    elif app.state.config.ENGINE == "comfyui":
+        try:
+            r = requests.get(url=f"{app.state.config.COMFYUI_BASE_URL}/object_info")
+            r.raise_for_status()
+            return True
+        except Exception as e:
+            app.state.config.ENABLED = False
+            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
+    else:
+        return True
+
+
 def set_image_model(model: str):
     app.state.config.MODEL = model
     if app.state.config.ENGINE in ["", "automatic1111"]:
diff --git a/src/lib/apis/images/index.ts b/src/lib/apis/images/index.ts
index ee5cda7f2..2e6510437 100644
--- a/src/lib/apis/images/index.ts
+++ b/src/lib/apis/images/index.ts
@@ -67,6 +67,38 @@ export const updateConfig = async (token: string = '', config: object) => {
 	return res;
 };
 
+export const verifyConfigUrl = async (token: string = '') => {
+	let error = null;
+
+	const res = await fetch(`${IMAGES_API_BASE_URL}/config/url/verify`, {
+		method: 'GET',
+		headers: {
+			Accept: 'application/json',
+			'Content-Type': 'application/json',
+			...(token && { authorization: `Bearer ${token}` })
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			if ('detail' in err) {
+				error = err.detail;
+			} else {
+				error = 'Server connection failed';
+			}
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const getImageGenerationConfig = async (token: string = '') => {
 	let error = null;
 
diff --git a/src/lib/components/admin/Settings/Images.svelte b/src/lib/components/admin/Settings/Images.svelte
index 1a792f764..58fcc690a 100644
--- a/src/lib/components/admin/Settings/Images.svelte
+++ b/src/lib/components/admin/Settings/Images.svelte
@@ -10,7 +10,8 @@
 		getImageGenerationConfig,
 		updateImageGenerationConfig,
 		getConfig,
-		updateConfig
+		updateConfig,
+		verifyConfigUrl
 	} from '$lib/apis/images';
 	import SensitiveInput from '$lib/components/common/SensitiveInput.svelte';
 	import Switch from '$lib/components/common/Switch.svelte';
@@ -255,8 +256,16 @@
 							<button
 								class="px-2.5 bg-gray-50 hover:bg-gray-100 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition"
 								type="button"
-								on:click={() => {
-									updateConfigHandler();
+								on:click={async () => {
+									await updateConfigHandler();
+									const res = await verifyConfigUrl(localStorage.token).catch((error) => {
+										toast.error(error);
+										return null;
+									});
+
+									if (res) {
+										toast.success($i18n.t('Server connection verified'));
+									}
 								}}
 							>
 								<svg
@@ -323,8 +332,16 @@
 							<button
 								class="px-2.5 bg-gray-50 hover:bg-gray-100 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition"
 								type="button"
-								on:click={() => {
-									updateConfigHandler();
+								on:click={async () => {
+									await updateConfigHandler();
+									const res = await verifyConfigUrl(localStorage.token).catch((error) => {
+										toast.error(error);
+										return null;
+									});
+
+									if (res) {
+										toast.success($i18n.t('Server connection verified'));
+									}
 								}}
 							>
 								<svg