main.py 13.8 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,
Self Denial's avatar
Self Denial committed
45
)
Timothy J. Baek's avatar
Timothy J. Baek committed
46
47


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

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
53
54
55
56
57
58
59
60
61
62

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

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

66
67
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
68

69
app.state.MODEL = IMAGE_GENERATION_MODEL
Timothy J. Baek's avatar
Timothy J. Baek committed
70
71


Timothy J. Baek's avatar
Timothy J. Baek committed
72
app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
73
74
app.state.COMFYUI_BASE_URL = COMFYUI_BASE_URL

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

76
77
app.state.IMAGE_SIZE = IMAGE_SIZE
app.state.IMAGE_STEPS = IMAGE_STEPS
Timothy J. Baek's avatar
Timothy J. Baek committed
78
79


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


Timothy J. Baek's avatar
Timothy J. Baek committed
85
86
87
88
89
90
91
92
93
94
class ConfigUpdateForm(BaseModel):
    engine: str
    enabled: bool


@app.post("/config/update")
async def update_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
    app.state.ENGINE = form_data.engine
    app.state.ENABLED = form_data.enabled
    return {"engine": app.state.ENGINE, "enabled": app.state.ENABLED}
Timothy J. Baek's avatar
Timothy J. Baek committed
95
96


Timothy J. Baek's avatar
Timothy J. Baek committed
97
98
99
class EngineUrlUpdateForm(BaseModel):
    AUTOMATIC1111_BASE_URL: Optional[str] = None
    COMFYUI_BASE_URL: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102


@app.get("/url")
Timothy J. Baek's avatar
Timothy J. Baek committed
103
104
105
106
107
async def get_engine_url(user=Depends(get_admin_user)):
    return {
        "AUTOMATIC1111_BASE_URL": app.state.AUTOMATIC1111_BASE_URL,
        "COMFYUI_BASE_URL": app.state.COMFYUI_BASE_URL,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
108
109
110


@app.post("/url/update")
Timothy J. Baek's avatar
Timothy J. Baek committed
111
112
async def update_engine_url(
    form_data: EngineUrlUpdateForm, user=Depends(get_admin_user)
Timothy J. Baek's avatar
Timothy J. Baek committed
113
):
Timothy J. Baek's avatar
Timothy J. Baek committed
114

Timothy J. Baek's avatar
Timothy J. Baek committed
115
    if form_data.AUTOMATIC1111_BASE_URL == None:
Timothy J. Baek's avatar
Timothy J. Baek committed
116
117
        app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
118
        url = form_data.AUTOMATIC1111_BASE_URL.strip("/")
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
121
122
123
        try:
            r = requests.head(url)
            app.state.AUTOMATIC1111_BASE_URL = url
        except Exception as e:
            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
124

Timothy J. Baek's avatar
Timothy J. Baek committed
125
126
127
128
    if form_data.COMFYUI_BASE_URL == None:
        app.state.COMFYUI_BASE_URL = COMFYUI_BASE_URL
    else:
        url = form_data.COMFYUI_BASE_URL.strip("/")
Timothy J. Baek's avatar
Timothy J. Baek committed
129
130
131
132
133
134

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

Timothy J. Baek's avatar
Timothy J. Baek committed
136
137
    return {
        "AUTOMATIC1111_BASE_URL": app.state.AUTOMATIC1111_BASE_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
138
        "COMFYUI_BASE_URL": app.state.COMFYUI_BASE_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
139
140
        "status": True,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
141
142


143
144
class OpenAIConfigUpdateForm(BaseModel):
    url: str
Timothy J. Baek's avatar
Timothy J. Baek committed
145
146
147
    key: str


148
149
150
151
152
153
@app.get("/openai/config")
async def get_openai_config(user=Depends(get_admin_user)):
    return {
        "OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL,
        "OPENAI_API_KEY": app.state.OPENAI_API_KEY,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
154
155


156
157
158
@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
159
160
161
162
):
    if form_data.key == "":
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)

163
    app.state.OPENAI_API_BASE_URL = form_data.url
Timothy J. Baek's avatar
Timothy J. Baek committed
164
    app.state.OPENAI_API_KEY = form_data.key
165

Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
    return {
        "status": True,
168
169
        "OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL,
        "OPENAI_API_KEY": app.state.OPENAI_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
170
171
172
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
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)."),
        )
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
198

199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223

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)."),
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
224
225


Timothy J. Baek's avatar
Timothy J. Baek committed
226
227
228
@app.get("/models")
def get_models(user=Depends(get_current_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
229
230
231
232
233
        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
234
235
236
237
238
239
240
241
242
243
244
245
        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
246
247
248
249
250
251
252
253
254
255
256
        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
257
    except Exception as e:
258
        app.state.ENABLED = False
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
259
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
260
261
262
263
264


@app.get("/models/default")
async def get_default_model(user=Depends(get_admin_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
        if app.state.ENGINE == "openai":
            return {"model": app.state.MODEL if app.state.MODEL else "dall-e-2"}
Timothy J. Baek's avatar
Timothy J. Baek committed
267
268
        elif app.state.ENGINE == "comfyui":
            return {"model": app.state.MODEL if app.state.MODEL else ""}
Timothy J. Baek's avatar
Timothy J. Baek committed
269
270
271
272
        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
273
    except Exception as e:
274
        app.state.ENABLED = False
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
275
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
276
277
278
279
280
281
282


class UpdateModelForm(BaseModel):
    model: str


def set_model_handler(model: str):
Timothy J. Baek's avatar
Timothy J. Baek committed
283
284
285
    if app.state.ENGINE == "openai":
        app.state.MODEL = model
        return app.state.MODEL
Timothy J. Baek's avatar
Timothy J. Baek committed
286
287
288
    if app.state.ENGINE == "comfyui":
        app.state.MODEL = model
        return app.state.MODEL
Timothy J. Baek's avatar
Timothy J. Baek committed
289
290
291
292
293
294
295
296
297
    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
298

Timothy J. Baek's avatar
Timothy J. Baek committed
299
        return options
Timothy J. Baek's avatar
Timothy J. Baek committed
300
301
302
303
304
305
306
307
308
309
310
311
312
313


@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
314
    size: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
315
316
317
    negative_prompt: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
318
319
def save_b64_image(b64_str):
    try:
320
321
        header, encoded = b64_str.split(",", 1)
        mime_type = header.split(";")[0]
Timothy J. Baek's avatar
Timothy J. Baek committed
322

323
        img_data = base64.b64decode(encoded)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
324

325
        image_id = str(uuid.uuid4())
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
326
327
328
329
        image_format = mimetypes.guess_extension(mime_type)

        image_filename = f"{image_id}{image_format}"
        file_path = IMAGE_CACHE_DIR / f"{image_filename}"
Timothy J. Baek's avatar
Timothy J. Baek committed
330
331
        with open(file_path, "wb") as f:
            f.write(img_data)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
332
        return image_filename
Timothy J. Baek's avatar
Timothy J. Baek committed
333
    except Exception as e:
334
        log.exception(f"Error saving image: {e}")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
335
        return None
Timothy J. Baek's avatar
Timothy J. Baek committed
336
337


Timothy J. Baek's avatar
Timothy J. Baek committed
338
339
340
341
342
def save_url_image(url):
    image_id = str(uuid.uuid4())
    try:
        r = requests.get(url)
        r.raise_for_status()
343
344
345
346
        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
347

348
349
350
            if not image_format:
                raise ValueError("Could not determine image type from MIME type")

Timothy J. Baek's avatar
Timothy J. Baek committed
351
352
353
            image_filename = f"{image_id}{image_format}"

            file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
354
355
356
            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
357
            return image_filename
358
359
        else:
            log.error(f"Url does not point to an image.")
Timothy J. Baek's avatar
Timothy J. Baek committed
360
            return None
Timothy J. Baek's avatar
Timothy J. Baek committed
361
362

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


Timothy J. Baek's avatar
Timothy J. Baek committed
367
368
369
370
371
372
@app.post("/generations")
def generate_image(
    form_data: GenerateImageForm,
    user=Depends(get_current_user),
):

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

Timothy J. Baek's avatar
Timothy J. Baek committed
375
    r = None
376
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
377
        if app.state.ENGINE == "openai":
378

Timothy J. Baek's avatar
Timothy J. Baek committed
379
380
381
            headers = {}
            headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
            headers["Content-Type"] = "application/json"
382

Timothy J. Baek's avatar
Timothy J. Baek committed
383
384
385
386
            data = {
                "model": app.state.MODEL if app.state.MODEL != "" else "dall-e-2",
                "prompt": form_data.prompt,
                "n": form_data.n,
Timothy J. Baek's avatar
Timothy J. Baek committed
387
                "size": form_data.size if form_data.size else app.state.IMAGE_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
388
389
                "response_format": "b64_json",
            }
390

Timothy J. Baek's avatar
Timothy J. Baek committed
391
            r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
392
                url=f"{app.state.OPENAI_API_BASE_URL}/images/generations",
Timothy J. Baek's avatar
Timothy J. Baek committed
393
394
395
                json=data,
                headers=headers,
            )
396

Timothy J. Baek's avatar
Timothy J. Baek committed
397
398
            r.raise_for_status()
            res = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
399

Timothy J. Baek's avatar
Timothy J. Baek committed
400
401
402
            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
403
404
                image_filename = save_b64_image(image["b64_json"])
                images.append({"url": f"/cache/image/generations/{image_filename}"})
Timothy J. Baek's avatar
Timothy J. Baek committed
405
406
407
408
409
410
411
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json")

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

            return images

Timothy J. Baek's avatar
Timothy J. Baek committed
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
        elif app.state.ENGINE == "comfyui":

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

            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

            data = ImageGenerationPayload(**data)

            res = comfyui_generate_image(
                app.state.MODEL,
                data,
                user.id,
                app.state.COMFYUI_BASE_URL,
            )
435
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
436
437
438
439

            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
440
441
442
                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
443
444
445
446

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

447
            log.debug(f"images: {images}")
Timothy J. Baek's avatar
Timothy J. Baek committed
448
            return images
Timothy J. Baek's avatar
Timothy J. Baek committed
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
        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,
            }

            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

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

            res = r.json()

473
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
474
475
476
477

            images = []

            for image in res["images"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
478
479
                image_filename = save_b64_image(image)
                images.append({"url": f"/cache/image/generations/{image_filename}"})
Timothy J. Baek's avatar
Timothy J. Baek committed
480
481
482
483
484
485
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json")

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

            return images
486
487

    except Exception as e:
488
489
490
491
492
493
494
        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))