diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index 39d3f96a..08c9a083 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -100,6 +100,32 @@ async def update_image_size( status_code=400, detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 512x512)."), ) + + +class ImageStepsUpdateForm(BaseModel): + steps: int + + +@app.get("/steps") +async def get_image_size(user=Depends(get_admin_user)): + return {"IMAGE_STEPS": app.state.IMAGE_STEPS} + + +@app.post("/steps/update") +async def update_image_size( + form_data: ImageStepsUpdateForm, user=Depends(get_admin_user) +): + if form_data.steps >= 0: + app.state.IMAGE_STEPS = form_data.steps + return { + "IMAGE_STEPS": app.state.IMAGE_STEPS, + "status": True, + } + else: + raise HTTPException( + status_code=400, + detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 50)."), + ) @app.get("/models") @@ -177,6 +203,9 @@ def generate_image( "height": height, } + if app.state.IMAGE_STEPS != None: + data["steps"] = app.state.IMAGE_STEPS + if form_data.negative_prompt != None: data["negative_prompt"] = form_data.negative_prompt diff --git a/src/lib/apis/images/index.ts b/src/lib/apis/images/index.ts index 205ee90a..c64bf2ba 100644 --- a/src/lib/apis/images/index.ts +++ b/src/lib/apis/images/index.ts @@ -198,6 +198,39 @@ export const updateImageSize = async (token: string = '', size: string) => { return res.IMAGE_SIZE; }; +export const updateImageSteps = async (token: string = '', steps: number) => { + let error = null; + + const res = await fetch(`${IMAGES_API_BASE_URL}/steps/update`, { + method: 'POST', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + ...(token && { authorization: `Bearer ${token}` }) + }, + body: JSON.stringify({ steps }) + }) + .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_STEPS; +}; + 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 94cb3213..2f6daa10 100644 --- a/src/lib/components/chat/Settings/Images.svelte +++ b/src/lib/components/chat/Settings/Images.svelte @@ -12,7 +12,8 @@ toggleImageGenerationEnabledStatus, updateAUTOMATIC1111Url, updateDefaultDiffusionModel, - updateImageSize + updateImageSize, + updateImageSteps } from '$lib/apis/images'; import { getBackendConfig } from '$lib/apis'; const dispatch = createEventDispatcher(); @@ -28,6 +29,7 @@ let models = []; let imageSize = ''; + let steps = 50; const getModels = async () => { models = await getDiffusionModels(localStorage.token).catch((error) => { @@ -98,6 +100,10 @@ toast.error(error); return null; }); + await updateImageSteps(localStorage.token, steps).catch((error) => { + toast.error(error); + return null; + }); dispatch('save'); loading = false; @@ -210,6 +216,19 @@ + +
+
Set Steps
+
+
+ +
+
+
{/if}