From 7ec4c07bf9c3607397ddbc5a53832adf255def5d Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Thu, 22 Feb 2024 19:32:36 -0800 Subject: [PATCH 1/2] feat: image size param --- backend/apps/images/main.py | 32 ++++++++- backend/constants.py | 3 + src/lib/apis/images/index.ts | 67 +++++++++++++++++++ .../components/chat/Settings/Images.svelte | 29 ++++++-- 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index 998af3ddb..39d3f96aa 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -1,4 +1,4 @@ -import os +import re import requests from fastapi import ( FastAPI, @@ -34,6 +34,7 @@ app.add_middleware( app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL app.state.ENABLED = app.state.AUTOMATIC1111_BASE_URL != "" +app.state.IMAGE_SIZE = "512x512" @app.get("/enabled", response_model=bool) @@ -74,6 +75,33 @@ async def update_openai_url(form_data: UrlUpdateForm, user=Depends(get_admin_use } +class ImageSizeUpdateForm(BaseModel): + size: str + + +@app.get("/size") +async def get_image_size(user=Depends(get_admin_user)): + return {"IMAGE_SIZE": app.state.IMAGE_SIZE} + + +@app.post("/size/update") +async def update_image_size( + form_data: ImageSizeUpdateForm, user=Depends(get_admin_user) +): + pattern = r"^\d+x\d+$" # Regular expression pattern + if re.match(pattern, form_data.size): + app.state.IMAGE_SIZE = form_data.size + return { + "IMAGE_SIZE": app.state.IMAGE_SIZE, + "status": True, + } + else: + raise HTTPException( + status_code=400, + detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 512x512)."), + ) + + @app.get("/models") def get_models(user=Depends(get_current_user)): try: @@ -140,7 +168,7 @@ def generate_image( if form_data.model: set_model_handler(form_data.model) - width, height = tuple(map(int, form_data.size.split("x"))) + width, height = tuple(map(int, app.state.IMAGE_SIZE.split("x"))) data = { "prompt": form_data.prompt, diff --git a/backend/constants.py b/backend/constants.py index 580db9c5f..cb802c7f4 100644 --- a/backend/constants.py +++ b/backend/constants.py @@ -44,3 +44,6 @@ class ERROR_MESSAGES(str, Enum): MALICIOUS = "Unusual activities detected, please try again in a few minutes." PANDOC_NOT_INSTALLED = "Pandoc is not installed on the server. Please contact your administrator for assistance." + INCORRECT_FORMAT = ( + lambda err="": f"Invalid format. Please use the correct format{err if err else ''}" + ) diff --git a/src/lib/apis/images/index.ts b/src/lib/apis/images/index.ts index b25499d64..205ee90ac 100644 --- a/src/lib/apis/images/index.ts +++ b/src/lib/apis/images/index.ts @@ -131,6 +131,73 @@ export const updateAUTOMATIC1111Url = async (token: string = '', url: string) => return res.AUTOMATIC1111_BASE_URL; }; +export const getImageSize = async (token: string = '') => { + let error = null; + + const res = await fetch(`${IMAGES_API_BASE_URL}/size`, { + 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.IMAGE_SIZE; +}; + +export const updateImageSize = async (token: string = '', size: string) => { + let error = null; + + const res = await fetch(`${IMAGES_API_BASE_URL}/size/update`, { + method: 'POST', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + ...(token && { authorization: `Bearer ${token}` }) + }, + body: JSON.stringify({ + size: size + }) + }) + .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.IMAGE_SIZE; +}; + export const getDiffusionModels = async (token: string = '') => { let error = null; diff --git a/src/lib/components/chat/Settings/Images.svelte b/src/lib/components/chat/Settings/Images.svelte index 9b4b60cdd..e6809c767 100644 --- a/src/lib/components/chat/Settings/Images.svelte +++ b/src/lib/components/chat/Settings/Images.svelte @@ -8,9 +8,11 @@ getDefaultDiffusionModel, getDiffusionModels, getImageGenerationEnabledStatus, + getImageSize, toggleImageGenerationEnabledStatus, updateAUTOMATIC1111Url, - updateDefaultDiffusionModel + updateDefaultDiffusionModel, + updateImageSize } from '$lib/apis/images'; import { getBackendConfig } from '$lib/apis'; const dispatch = createEventDispatcher(); @@ -25,6 +27,8 @@ let selectedModel = ''; let models = []; + let imageSize = ''; + const getModels = async () => { models = await getDiffusionModels(localStorage.token).catch((error) => { toast.error(error); @@ -53,7 +57,6 @@ AUTOMATIC1111_BASE_URL = await getAUTOMATIC1111Url(localStorage.token); } }; - const toggleImageGeneration = async () => { if (AUTOMATIC1111_BASE_URL) { enableImageGeneration = await toggleImageGenerationEnabledStatus(localStorage.token).catch( @@ -79,6 +82,7 @@ AUTOMATIC1111_BASE_URL = await getAUTOMATIC1111Url(localStorage.token); if (enableImageGeneration && AUTOMATIC1111_BASE_URL) { + imageSize = await getImageSize(localStorage.token); getModels(); } } @@ -89,13 +93,17 @@ class="flex flex-col h-full justify-between space-y-3 text-sm" on:submit|preventDefault={async () => { loading = true; - const res = await updateDefaultDiffusionModel(localStorage.token, selectedModel); + await updateDefaultDiffusionModel(localStorage.token, selectedModel); + await updateImageSize(localStorage.token, imageSize).catch((error) => { + toast.error(error); + return null; + }); dispatch('save'); loading = false; }} > -
+
Image Settings
@@ -168,6 +176,19 @@ {#if enableImageGeneration}
+
+
Set Image Size
+
+
+ +
+
+
+
Set default model
From e612feec9925898aa94649e820e8c872939cf759 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Thu, 22 Feb 2024 19:33:36 -0800 Subject: [PATCH 2/2] Update package.json --- package.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/package.json b/package.json index 352130769..332362747 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "open-webui", - "version": "0.1.0-101", + "version": "0.1.102", "private": true, "scripts": { "dev": "vite dev --host", @@ -52,4 +52,4 @@ "tippy.js": "^6.3.7", "uuid": "^9.0.1" } -} +} \ No newline at end of file