main.py 14.2 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:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
320
        image_id = str(uuid.uuid4())
Timothy J. Baek's avatar
Timothy J. Baek committed
321

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
322
323
324
        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
325

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
            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
344

Timothy J. Baek's avatar
Timothy J. Baek committed
345
    except Exception as e:
346
        log.exception(f"Error saving image: {e}")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
347
        return None
Timothy J. Baek's avatar
Timothy J. Baek committed
348
349


Timothy J. Baek's avatar
Timothy J. Baek committed
350
351
352
353
354
def save_url_image(url):
    image_id = str(uuid.uuid4())
    try:
        r = requests.get(url)
        r.raise_for_status()
355
356
357
358
        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
359

360
361
362
            if not image_format:
                raise ValueError("Could not determine image type from MIME type")

Timothy J. Baek's avatar
Timothy J. Baek committed
363
364
365
            image_filename = f"{image_id}{image_format}"

            file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
366
367
368
            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
369
            return image_filename
370
371
        else:
            log.error(f"Url does not point to an image.")
Timothy J. Baek's avatar
Timothy J. Baek committed
372
            return None
Timothy J. Baek's avatar
Timothy J. Baek committed
373
374

    except Exception as e:
375
        log.exception(f"Error saving image: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
376
        return None
Timothy J. Baek's avatar
Timothy J. Baek committed
377
378


Timothy J. Baek's avatar
Timothy J. Baek committed
379
380
381
382
383
384
@app.post("/generations")
def generate_image(
    form_data: GenerateImageForm,
    user=Depends(get_current_user),
):

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

Timothy J. Baek's avatar
Timothy J. Baek committed
387
    r = None
388
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
389
        if app.state.ENGINE == "openai":
390

Timothy J. Baek's avatar
Timothy J. Baek committed
391
392
393
            headers = {}
            headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
            headers["Content-Type"] = "application/json"
394

Timothy J. Baek's avatar
Timothy J. Baek committed
395
396
397
398
            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
399
                "size": form_data.size if form_data.size else app.state.IMAGE_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
400
401
                "response_format": "b64_json",
            }
402

Timothy J. Baek's avatar
Timothy J. Baek committed
403
            r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
404
                url=f"{app.state.OPENAI_API_BASE_URL}/images/generations",
Timothy J. Baek's avatar
Timothy J. Baek committed
405
406
407
                json=data,
                headers=headers,
            )
408

Timothy J. Baek's avatar
Timothy J. Baek committed
409
410
            r.raise_for_status()
            res = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
411

Timothy J. Baek's avatar
Timothy J. Baek committed
412
413
414
            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
415
416
                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
417
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
Timothy J. Baek's avatar
Timothy J. Baek committed
418
419
420
421
422
423

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

            return images

Timothy J. Baek's avatar
Timothy J. Baek committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
        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,
            )
447
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
448
449
450
451

            images = []

            for image in res["data"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
452
453
454
                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
455
456
457
458

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

459
            log.debug(f"images: {images}")
Timothy J. Baek's avatar
Timothy J. Baek committed
460
            return images
Timothy J. Baek's avatar
Timothy J. Baek committed
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
        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()

485
            log.debug(f"res: {res}")
Timothy J. Baek's avatar
Timothy J. Baek committed
486
487
488
489

            images = []

            for image in res["images"]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
490
491
                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
492
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
Timothy J. Baek's avatar
Timothy J. Baek committed
493
494
495
496
497

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

            return images
498
499

    except Exception as e:
500
501
502
503
504
505
506
        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))