main.py 14.9 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
import re
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import requests
from fastapi import (
    FastAPI,
    Request,
    Depends,
    HTTPException,
    status,
    UploadFile,
    File,
    Form,
)
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel

from constants import ERROR_MESSAGES
from utils.utils import (
    get_current_user,
    get_admin_user,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
21
22

from apps.images.utils.comfyui import ImageGenerationPayload, comfyui_generate_image
Timothy J. Baek's avatar
Timothy J. Baek committed
23
24
25
from utils.misc import calculate_sha256
from typing import Optional
from pydantic import BaseModel
Timothy J. Baek's avatar
Timothy J. Baek committed
26
from pathlib import Path
27
import mimetypes
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
30
import uuid
import base64
import json
31
import logging
Timothy J. Baek's avatar
Timothy J. Baek committed
32

Self Denial's avatar
Self Denial committed
33
34
35
from config import (
    SRC_LOG_LEVELS,
    CACHE_DIR,
36
    IMAGE_GENERATION_ENGINE,
Self Denial's avatar
Self Denial committed
37
    ENABLE_IMAGE_GENERATION,
Self Denial's avatar
Self Denial committed
38
39
    AUTOMATIC1111_BASE_URL,
    COMFYUI_BASE_URL,
40
41
    IMAGES_OPENAI_API_BASE_URL,
    IMAGES_OPENAI_API_KEY,
42
    IMAGE_GENERATION_MODEL,
43
44
    IMAGE_SIZE,
    IMAGE_STEPS,
45
46
    config_get,
    config_set,
Self Denial's avatar
Self Denial committed
47
)
Timothy J. Baek's avatar
Timothy J. Baek committed
48
49


50
51
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["IMAGES"])
Timothy J. Baek's avatar
Timothy J. Baek committed
52
53
54

IMAGE_CACHE_DIR = Path(CACHE_DIR).joinpath("./image/generations/")
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
55
56
57
58
59
60
61
62
63
64

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

65
app.state.ENGINE = IMAGE_GENERATION_ENGINE
Self Denial's avatar
Self Denial committed
66
app.state.ENABLED = ENABLE_IMAGE_GENERATION
Timothy J. Baek's avatar
Timothy J. Baek committed
67

68
69
app.state.OPENAI_API_BASE_URL = IMAGES_OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = IMAGES_OPENAI_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
70

71
app.state.MODEL = IMAGE_GENERATION_MODEL
Timothy J. Baek's avatar
Timothy J. Baek committed
72
73


Timothy J. Baek's avatar
Timothy J. Baek committed
74
app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
75
76
app.state.COMFYUI_BASE_URL = COMFYUI_BASE_URL

Timothy J. Baek's avatar
Timothy J. Baek committed
77

78
79
app.state.IMAGE_SIZE = IMAGE_SIZE
app.state.IMAGE_STEPS = IMAGE_STEPS
Timothy J. Baek's avatar
Timothy J. Baek committed
80
81


Timothy J. Baek's avatar
Timothy J. Baek committed
82
83
@app.get("/config")
async def get_config(request: Request, user=Depends(get_admin_user)):
84
85
86
87
    return {
        "engine": config_get(app.state.ENGINE),
        "enabled": config_get(app.state.ENABLED),
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
88
89


Timothy J. Baek's avatar
Timothy J. Baek committed
90
91
92
93
94
95
96
class ConfigUpdateForm(BaseModel):
    engine: str
    enabled: bool


@app.post("/config/update")
async def update_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
97
98
99
100
101
102
    config_set(app.state.ENGINE, form_data.engine)
    config_set(app.state.ENABLED, form_data.enabled)
    return {
        "engine": config_get(app.state.ENGINE),
        "enabled": config_get(app.state.ENABLED),
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
103
104


Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
107
class EngineUrlUpdateForm(BaseModel):
    AUTOMATIC1111_BASE_URL: Optional[str] = None
    COMFYUI_BASE_URL: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
108
109
110


@app.get("/url")
Timothy J. Baek's avatar
Timothy J. Baek committed
111
112
async def get_engine_url(user=Depends(get_admin_user)):
    return {
113
114
        "AUTOMATIC1111_BASE_URL": config_get(app.state.AUTOMATIC1111_BASE_URL),
        "COMFYUI_BASE_URL": config_get(app.state.COMFYUI_BASE_URL),
Timothy J. Baek's avatar
Timothy J. Baek committed
115
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
116
117
118


@app.post("/url/update")
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
async def update_engine_url(
    form_data: EngineUrlUpdateForm, user=Depends(get_admin_user)
Timothy J. Baek's avatar
Timothy J. Baek committed
121
):
Timothy J. Baek's avatar
Timothy J. Baek committed
122

Timothy J. Baek's avatar
Timothy J. Baek committed
123
    if form_data.AUTOMATIC1111_BASE_URL == None:
124
        config_set(app.state.AUTOMATIC1111_BASE_URL, config_get(AUTOMATIC1111_BASE_URL))
Timothy J. Baek's avatar
Timothy J. Baek committed
125
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
126
        url = form_data.AUTOMATIC1111_BASE_URL.strip("/")
Timothy J. Baek's avatar
Timothy J. Baek committed
127
128
        try:
            r = requests.head(url)
129
            config_set(app.state.AUTOMATIC1111_BASE_URL, url)
Timothy J. Baek's avatar
Timothy J. Baek committed
130
131
        except Exception as e:
            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
132

Timothy J. Baek's avatar
Timothy J. Baek committed
133
    if form_data.COMFYUI_BASE_URL == None:
134
        config_set(app.state.COMFYUI_BASE_URL, COMFYUI_BASE_URL)
Timothy J. Baek's avatar
Timothy J. Baek committed
135
136
    else:
        url = form_data.COMFYUI_BASE_URL.strip("/")
Timothy J. Baek's avatar
Timothy J. Baek committed
137
138
139

        try:
            r = requests.head(url)
140
            config_set(app.state.COMFYUI_BASE_URL, url)
Timothy J. Baek's avatar
Timothy J. Baek committed
141
142
        except Exception as e:
            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
143

Timothy J. Baek's avatar
Timothy J. Baek committed
144
    return {
145
146
        "AUTOMATIC1111_BASE_URL": config_get(app.state.AUTOMATIC1111_BASE_URL),
        "COMFYUI_BASE_URL": config_get(app.state.COMFYUI_BASE_URL),
Timothy J. Baek's avatar
Timothy J. Baek committed
147
148
        "status": True,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
149
150


151
152
class OpenAIConfigUpdateForm(BaseModel):
    url: str
Timothy J. Baek's avatar
Timothy J. Baek committed
153
154
155
    key: str


156
157
158
@app.get("/openai/config")
async def get_openai_config(user=Depends(get_admin_user)):
    return {
159
160
        "OPENAI_API_BASE_URL": config_get(app.state.OPENAI_API_BASE_URL),
        "OPENAI_API_KEY": config_get(app.state.OPENAI_API_KEY),
161
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
162
163


164
165
166
@app.post("/openai/config/update")
async def update_openai_config(
    form_data: OpenAIConfigUpdateForm, user=Depends(get_admin_user)
Timothy J. Baek's avatar
Timothy J. Baek committed
167
168
169
170
):
    if form_data.key == "":
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)

171
172
    config_set(app.state.OPENAI_API_BASE_URL, form_data.url)
    config_set(app.state.OPENAI_API_KEY, form_data.key)
173

Timothy J. Baek's avatar
Timothy J. Baek committed
174
175
    return {
        "status": True,
176
177
        "OPENAI_API_BASE_URL": config_get(app.state.OPENAI_API_BASE_URL),
        "OPENAI_API_KEY": config_get(app.state.OPENAI_API_KEY),
Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
180
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
181
182
183
184
185
186
class ImageSizeUpdateForm(BaseModel):
    size: str


@app.get("/size")
async def get_image_size(user=Depends(get_admin_user)):
187
    return {"IMAGE_SIZE": config_get(app.state.IMAGE_SIZE)}
Timothy J. Baek's avatar
Timothy J. Baek committed
188
189
190
191
192
193
194
195


@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):
196
        config_set(app.state.IMAGE_SIZE, form_data.size)
Timothy J. Baek's avatar
Timothy J. Baek committed
197
        return {
198
            "IMAGE_SIZE": config_get(app.state.IMAGE_SIZE),
Timothy J. Baek's avatar
Timothy J. Baek committed
199
200
201
202
203
204
205
            "status": True,
        }
    else:
        raise HTTPException(
            status_code=400,
            detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 512x512)."),
        )
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
206

207
208
209
210
211
212
213

class ImageStepsUpdateForm(BaseModel):
    steps: int


@app.get("/steps")
async def get_image_size(user=Depends(get_admin_user)):
214
    return {"IMAGE_STEPS": config_get(app.state.IMAGE_STEPS)}
215
216
217
218
219
220
221


@app.post("/steps/update")
async def update_image_size(
    form_data: ImageStepsUpdateForm, user=Depends(get_admin_user)
):
    if form_data.steps >= 0:
222
        config_set(app.state.IMAGE_STEPS, form_data.steps)
223
        return {
224
            "IMAGE_STEPS": config_get(app.state.IMAGE_STEPS),
225
226
227
228
229
230
231
            "status": True,
        }
    else:
        raise HTTPException(
            status_code=400,
            detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 50)."),
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
232
233


Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
236
@app.get("/models")
def get_models(user=Depends(get_current_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
237
238
239
240
241
        if app.state.ENGINE == "openai":
            return [
                {"id": "dall-e-2", "name": "DALL·E 2"},
                {"id": "dall-e-3", "name": "DALL·E 3"},
            ]
Timothy J. Baek's avatar
Timothy J. Baek committed
242
243
244
245
246
247
248
249
250
251
252
253
        elif app.state.ENGINE == "comfyui":

            r = requests.get(url=f"{app.state.COMFYUI_BASE_URL}/object_info")
            info = r.json()

            return list(
                map(
                    lambda model: {"id": model, "name": model},
                    info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0],
                )
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
254
255
256
257
258
259
260
261
262
263
264
        else:
            r = requests.get(
                url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models"
            )
            models = r.json()
            return list(
                map(
                    lambda model: {"id": model["title"], "name": model["model_name"]},
                    models,
                )
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
265
    except Exception as e:
266
        app.state.ENABLED = False
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
267
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
268
269
270
271
272


@app.get("/models/default")
async def get_default_model(user=Depends(get_admin_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
273
        if app.state.ENGINE == "openai":
274
275
276
277
278
279
280
            return {
                "model": (
                    config_get(app.state.MODEL)
                    if config_get(app.state.MODEL)
                    else "dall-e-2"
                )
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
281
        elif app.state.ENGINE == "comfyui":
282
283
284
285
286
            return {
                "model": (
                    config_get(app.state.MODEL) if config_get(app.state.MODEL) else ""
                )
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
287
288
289
290
        else:
            r = requests.get(url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/options")
            options = r.json()
            return {"model": options["sd_model_checkpoint"]}
Timothy J. Baek's avatar
Timothy J. Baek committed
291
    except Exception as e:
292
        config_set(app.state.ENABLED, False)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
293
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
296
297
298
299
300


class UpdateModelForm(BaseModel):
    model: str


def set_model_handler(model: str):
301
302
303
    if app.state.ENGINE in ["openai", "comfyui"]:
        config_set(app.state.MODEL, model)
        return config_get(app.state.MODEL)
Timothy J. Baek's avatar
Timothy J. Baek committed
304
305
306
307
308
309
310
311
312
    else:
        r = requests.get(url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/options")
        options = r.json()

        if model != options["sd_model_checkpoint"]:
            options["sd_model_checkpoint"] = model
            r = requests.post(
                url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/options", json=options
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
313

Timothy J. Baek's avatar
Timothy J. Baek committed
314
        return options
Timothy J. Baek's avatar
Timothy J. Baek committed
315
316
317
318
319
320
321
322
323
324
325
326
327
328


@app.post("/models/default/update")
def update_default_model(
    form_data: UpdateModelForm,
    user=Depends(get_current_user),
):
    return set_model_handler(form_data.model)


class GenerateImageForm(BaseModel):
    model: Optional[str] = None
    prompt: str
    n: int = 1
Timothy J. Baek's avatar
Timothy J. Baek committed
329
    size: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
330
331
332
    negative_prompt: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
333
334
def save_b64_image(b64_str):
    try:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
335
        image_id = str(uuid.uuid4())
Timothy J. Baek's avatar
Timothy J. Baek committed
336

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
337
338
339
        if "," in b64_str:
            header, encoded = b64_str.split(",", 1)
            mime_type = header.split(";")[0]
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
340

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
            img_data = base64.b64decode(encoded)
            image_format = mimetypes.guess_extension(mime_type)

            image_filename = f"{image_id}{image_format}"
            file_path = IMAGE_CACHE_DIR / f"{image_filename}"
            with open(file_path, "wb") as f:
                f.write(img_data)
            return image_filename
        else:
            image_filename = f"{image_id}.png"
            file_path = IMAGE_CACHE_DIR.joinpath(image_filename)

            img_data = base64.b64decode(b64_str)

            # Write the image data to a file
            with open(file_path, "wb") as f:
                f.write(img_data)
            return image_filename
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
359

Timothy J. Baek's avatar
Timothy J. Baek committed
360
    except Exception as e:
361
        log.exception(f"Error saving image: {e}")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
362
        return None
Timothy J. Baek's avatar
Timothy J. Baek committed
363
364


Timothy J. Baek's avatar
Timothy J. Baek committed
365
366
367
368
369
def save_url_image(url):
    image_id = str(uuid.uuid4())
    try:
        r = requests.get(url)
        r.raise_for_status()
370
371
372
373
        if r.headers["content-type"].split("/")[0] == "image":

            mime_type = r.headers["content-type"]
            image_format = mimetypes.guess_extension(mime_type)
Timothy J. Baek's avatar
Timothy J. Baek committed
374

375
376
377
            if not image_format:
                raise ValueError("Could not determine image type from MIME type")

Timothy J. Baek's avatar
Timothy J. Baek committed
378
379
380
            image_filename = f"{image_id}{image_format}"

            file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
381
382
383
            with open(file_path, "wb") as image_file:
                for chunk in r.iter_content(chunk_size=8192):
                    image_file.write(chunk)
Timothy J. Baek's avatar
Timothy J. Baek committed
384
            return image_filename
385
386
        else:
            log.error(f"Url does not point to an image.")
Timothy J. Baek's avatar
Timothy J. Baek committed
387
            return None
Timothy J. Baek's avatar
Timothy J. Baek committed
388
389

    except Exception as e:
390
        log.exception(f"Error saving image: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
391
        return None
Timothy J. Baek's avatar
Timothy J. Baek committed
392
393


Timothy J. Baek's avatar
Timothy J. Baek committed
394
395
396
397
398
399
@app.post("/generations")
def generate_image(
    form_data: GenerateImageForm,
    user=Depends(get_current_user),
):

400
    width, height = tuple(map(int, config_get(app.state.IMAGE_SIZE).split("x")))
Timothy J. Baek's avatar
Timothy J. Baek committed
401

Timothy J. Baek's avatar
Timothy J. Baek committed
402
    r = None
403
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
404
        if app.state.ENGINE == "openai":
405

Timothy J. Baek's avatar
Timothy J. Baek committed
406
407
408
            headers = {}
            headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
            headers["Content-Type"] = "application/json"
409

Timothy J. Baek's avatar
Timothy J. Baek committed
410
411
412
413
            data = {
                "model": app.state.MODEL if app.state.MODEL != "" else "dall-e-2",
                "prompt": form_data.prompt,
                "n": form_data.n,
414
415
416
417
418
                "size": (
                    form_data.size
                    if form_data.size
                    else config_get(app.state.IMAGE_SIZE)
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
419
420
                "response_format": "b64_json",
            }
421

Timothy J. Baek's avatar
Timothy J. Baek committed
422
            r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
423
                url=f"{app.state.OPENAI_API_BASE_URL}/images/generations",
Timothy J. Baek's avatar
Timothy J. Baek committed
424
425
426
                json=data,
                headers=headers,
            )
427

Timothy J. Baek's avatar
Timothy J. Baek committed
428
429
            r.raise_for_status()
            res = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
430

Timothy J. Baek's avatar
Timothy J. Baek committed
431
432
433
            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
434
435
                image_filename = save_b64_image(image["b64_json"])
                images.append({"url": f"/cache/image/generations/{image_filename}"})
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
436
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
Timothy J. Baek's avatar
Timothy J. Baek committed
437
438
439
440
441
442

                with open(file_body_path, "w") as f:
                    json.dump(data, f)

            return images

Timothy J. Baek's avatar
Timothy J. Baek committed
443
444
445
446
447
448
449
450
451
        elif app.state.ENGINE == "comfyui":

            data = {
                "prompt": form_data.prompt,
                "width": width,
                "height": height,
                "n": form_data.n,
            }

452
453
            if config_get(app.state.IMAGE_STEPS) is not None:
                data["steps"] = config_get(app.state.IMAGE_STEPS)
Timothy J. Baek's avatar
Timothy J. Baek committed
454

455
            if form_data.negative_prompt is not None:
Timothy J. Baek's avatar
Timothy J. Baek committed
456
457
458
459
460
                data["negative_prompt"] = form_data.negative_prompt

            data = ImageGenerationPayload(**data)

            res = comfyui_generate_image(
461
                config_get(app.state.MODEL),
Timothy J. Baek's avatar
Timothy J. Baek committed
462
463
                data,
                user.id,
464
                config_get(app.state.COMFYUI_BASE_URL),
Timothy J. Baek's avatar
Timothy J. Baek committed
465
            )
466
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
467
468
469
470

            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
471
472
473
                image_filename = save_url_image(image["url"])
                images.append({"url": f"/cache/image/generations/{image_filename}"})
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
Timothy J. Baek's avatar
Timothy J. Baek committed
474
475
476
477

                with open(file_body_path, "w") as f:
                    json.dump(data.model_dump(exclude_none=True), f)

478
            log.debug(f"images: {images}")
Timothy J. Baek's avatar
Timothy J. Baek committed
479
            return images
Timothy J. Baek's avatar
Timothy J. Baek committed
480
481
482
483
484
485
486
487
488
489
490
        else:
            if form_data.model:
                set_model_handler(form_data.model)

            data = {
                "prompt": form_data.prompt,
                "batch_size": form_data.n,
                "width": width,
                "height": height,
            }

491
492
            if config_get(app.state.IMAGE_STEPS) is not None:
                data["steps"] = config_get(app.state.IMAGE_STEPS)
Timothy J. Baek's avatar
Timothy J. Baek committed
493

494
            if form_data.negative_prompt is not None:
Timothy J. Baek's avatar
Timothy J. Baek committed
495
496
497
498
499
500
501
502
503
                data["negative_prompt"] = form_data.negative_prompt

            r = requests.post(
                url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img",
                json=data,
            )

            res = r.json()

504
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
505
506
507
508

            images = []

            for image in res["images"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
509
510
                image_filename = save_b64_image(image)
                images.append({"url": f"/cache/image/generations/{image_filename}"})
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
511
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
Timothy J. Baek's avatar
Timothy J. Baek committed
512
513
514
515
516

                with open(file_body_path, "w") as f:
                    json.dump({**data, "info": res["info"]}, f)

            return images
517
518

    except Exception as e:
519
520
521
522
523
524
525
        error = e

        if r != None:
            data = r.json()
            if "error" in data:
                error = data["error"]["message"]
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error))