main.py 9.55 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
21
22
23
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,
)
from utils.misc import calculate_sha256
from typing import Optional
from pydantic import BaseModel
Timothy J. Baek's avatar
Timothy J. Baek committed
24
25
26
27
28
29
30
31
32
33
from pathlib import Path
import uuid
import base64
import json

from config import CACHE_DIR, AUTOMATIC1111_BASE_URL


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
34
35
36
37
38
39
40
41
42
43

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

Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
46
47
48
49
50
app.state.ENGINE = ""
app.state.ENABLED = False

app.state.OPENAI_API_KEY = ""
app.state.MODEL = ""


Timothy J. Baek's avatar
Timothy J. Baek committed
51
app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
52

Timothy J. Baek's avatar
Timothy J. Baek committed
53
app.state.IMAGE_SIZE = "512x512"
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
54
app.state.IMAGE_STEPS = 50
Timothy J. Baek's avatar
Timothy J. Baek committed
55
56


Timothy J. Baek's avatar
Timothy J. Baek committed
57
58
59
@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
60
61


Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
65
66
67
68
69
70
71
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
72
73
74
75
76
77
78


class UrlUpdateForm(BaseModel):
    url: str


@app.get("/url")
Timothy J. Baek's avatar
Timothy J. Baek committed
79
async def get_automatic1111_url(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
80
81
82
83
    return {"AUTOMATIC1111_BASE_URL": app.state.AUTOMATIC1111_BASE_URL}


@app.post("/url/update")
Timothy J. Baek's avatar
Timothy J. Baek committed
84
85
86
async def update_automatic1111_url(
    form_data: UrlUpdateForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
89
90

    if form_data.url == "":
        app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
91
92
93
94
95
96
        url = form_data.url.strip("/")
        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
97
98
99
100
101

    return {
        "AUTOMATIC1111_BASE_URL": app.state.AUTOMATIC1111_BASE_URL,
        "status": True,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
102
103


Timothy J. Baek's avatar
Timothy J. Baek committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
class OpenAIKeyUpdateForm(BaseModel):
    key: str


@app.get("/key")
async def get_openai_key(user=Depends(get_admin_user)):
    return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}


@app.post("/key/update")
async def update_openai_key(
    form_data: OpenAIKeyUpdateForm, user=Depends(get_admin_user)
):

    if form_data.key == "":
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)

    app.state.OPENAI_API_KEY = form_data.key
    return {
        "OPENAI_API_KEY": app.state.OPENAI_API_KEY,
        "status": True,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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
153

154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

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
179
180


Timothy J. Baek's avatar
Timothy J. Baek committed
181
182
183
@app.get("/models")
def get_models(user=Depends(get_current_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        if app.state.ENGINE == "openai":
            return [
                {"id": "dall-e-2", "name": "DALL·E 2"},
                {"id": "dall-e-3", "name": "DALL·E 3"},
            ]
        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
200
    except Exception as e:
201
        app.state.ENABLED = False
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
202
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
203
204
205
206
207


@app.get("/models/default")
async def get_default_model(user=Depends(get_admin_user)):
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
208
209
210
211
212
213
        if app.state.ENGINE == "openai":
            return {"model": app.state.MODEL if app.state.MODEL else "dall-e-2"}
        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
214
    except Exception as e:
215
        app.state.ENABLED = False
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
216
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
Timothy J. Baek's avatar
Timothy J. Baek committed
217
218
219
220
221
222
223
224


class UpdateModelForm(BaseModel):
    model: str


def set_model_handler(model: str):

Timothy J. Baek's avatar
Timothy J. Baek committed
225
226
227
228
229
230
231
232
233
234
235
236
    if app.state.ENGINE == "openai":
        app.state.MODEL = model
        return app.state.MODEL
    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
237

Timothy J. Baek's avatar
Timothy J. Baek committed
238
        return options
Timothy J. Baek's avatar
Timothy J. Baek committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252


@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
253
    size: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
254
255
256
    negative_prompt: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def save_b64_image(b64_str):
    image_id = str(uuid.uuid4())
    file_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.png")

    try:
        # Split the base64 string to get the actual image data
        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_id
    except Exception as e:
        print(f"Error saving image: {e}")
        return None


Timothy J. Baek's avatar
Timothy J. Baek committed
275
276
277
278
279
280
@app.post("/generations")
def generate_image(
    form_data: GenerateImageForm,
    user=Depends(get_current_user),
):

Timothy J. Baek's avatar
Timothy J. Baek committed
281
    r = None
282
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
283
        if app.state.ENGINE == "openai":
284

Timothy J. Baek's avatar
Timothy J. Baek committed
285
286
287
            headers = {}
            headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
            headers["Content-Type"] = "application/json"
288

Timothy J. Baek's avatar
Timothy J. Baek committed
289
290
291
292
            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
293
                "size": form_data.size if form_data.size else app.state.IMAGE_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
                "response_format": "b64_json",
            }
296

Timothy J. Baek's avatar
Timothy J. Baek committed
297
298
299
300
301
            r = requests.post(
                url=f"https://api.openai.com/v1/images/generations",
                json=data,
                headers=headers,
            )
302

Timothy J. Baek's avatar
Timothy J. Baek committed
303
304
            r.raise_for_status()
            res = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
305

Timothy J. Baek's avatar
Timothy J. Baek committed
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
            images = []

            for image in res["data"]:
                image_id = save_b64_image(image["b64_json"])
                images.append({"url": f"/cache/image/generations/{image_id}.png"})
                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

        else:
            if form_data.model:
                set_model_handler(form_data.model)

            width, height = tuple(map(int, app.state.IMAGE_SIZE.split("x")))

            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()

            print(res)

            images = []

            for image in res["images"]:
                image_id = save_b64_image(image)
                images.append({"url": f"/cache/image/generations/{image_id}.png"})
                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
357
358

    except Exception as e:
359
360
361
362
363
364
365
        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))