main.py

import re
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 apps.images.utils.comfyui import ImageGenerationPayload, comfyui_generate_image
from utils.misc import calculate_sha256
from typing import Optional
from pydantic import BaseModel
from pathlib import Path
import mimetypes
import uuid
import base64
import json
import logging

from config import (
    SRC_LOG_LEVELS,
    CACHE_DIR,
    IMAGE_GENERATION_ENGINE,
    ENABLE_IMAGE_GENERATION,
    AUTOMATIC1111_BASE_URL,
    COMFYUI_BASE_URL,
    IMAGES_OPENAI_API_BASE_URL,
    IMAGES_OPENAI_API_KEY,
    IMAGE_GENERATION_MODEL,
    IMAGE_SIZE,
    IMAGE_STEPS,
    config_get,
    config_set,
)


log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["IMAGES"])

IMAGE_CACHE_DIR = Path(CACHE_DIR).joinpath("./image/generations/")
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)

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

app.state.ENGINE = IMAGE_GENERATION_ENGINE
app.state.ENABLED = ENABLE_IMAGE_GENERATION

app.state.OPENAI_API_BASE_URL = IMAGES_OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = IMAGES_OPENAI_API_KEY

app.state.MODEL = IMAGE_GENERATION_MODEL


app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
app.state.COMFYUI_BASE_URL = COMFYUI_BASE_URL


app.state.IMAGE_SIZE = IMAGE_SIZE
app.state.IMAGE_STEPS = IMAGE_STEPS


@app.get("/config")
async def get_config(request: Request, user=Depends(get_admin_user)):
    return {
        "engine": config_get(app.state.ENGINE),
        "enabled": config_get(app.state.ENABLED),
    }


class ConfigUpdateForm(BaseModel):
    engine: str
    enabled: bool


@app.post("/config/update")
async def update_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
    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),
    }


class EngineUrlUpdateForm(BaseModel):
    AUTOMATIC1111_BASE_URL: Optional[str] = None
    COMFYUI_BASE_URL: Optional[str] = None


@app.get("/url")
async def get_engine_url(user=Depends(get_admin_user)):
    return {
        "AUTOMATIC1111_BASE_URL": config_get(app.state.AUTOMATIC1111_BASE_URL),
        "COMFYUI_BASE_URL": config_get(app.state.COMFYUI_BASE_URL),
    }


@app.post("/url/update")
async def update_engine_url(
    form_data: EngineUrlUpdateForm, user=Depends(get_admin_user)
):

    if form_data.AUTOMATIC1111_BASE_URL == None:
        config_set(app.state.AUTOMATIC1111_BASE_URL, config_get(AUTOMATIC1111_BASE_URL))
    else:
        url = form_data.AUTOMATIC1111_BASE_URL.strip("/")
        try:
            r = requests.head(url)
            config_set(app.state.AUTOMATIC1111_BASE_URL, url)
        except Exception as e:
            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))

    if form_data.COMFYUI_BASE_URL == None:
        config_set(app.state.COMFYUI_BASE_URL, COMFYUI_BASE_URL)
    else:
        url = form_data.COMFYUI_BASE_URL.strip("/")

        try:
            r = requests.head(url)
            config_set(app.state.COMFYUI_BASE_URL, url)
        except Exception as e:
            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))

    return {
        "AUTOMATIC1111_BASE_URL": config_get(app.state.AUTOMATIC1111_BASE_URL),
        "COMFYUI_BASE_URL": config_get(app.state.COMFYUI_BASE_URL),
        "status": True,
    }


class OpenAIConfigUpdateForm(BaseModel):
    url: str
    key: str


@app.get("/openai/config")
async def get_openai_config(user=Depends(get_admin_user)):
    return {
        "OPENAI_API_BASE_URL": config_get(app.state.OPENAI_API_BASE_URL),
        "OPENAI_API_KEY": config_get(app.state.OPENAI_API_KEY),
    }


@app.post("/openai/config/update")
async def update_openai_config(
    form_data: OpenAIConfigUpdateForm, user=Depends(get_admin_user)
):
    if form_data.key == "":
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)

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

    return {
        "status": True,
        "OPENAI_API_BASE_URL": config_get(app.state.OPENAI_API_BASE_URL),
        "OPENAI_API_KEY": config_get(app.state.OPENAI_API_KEY),
    }


class ImageSizeUpdateForm(BaseModel):
    size: str


@app.get("/size")
async def get_image_size(user=Depends(get_admin_user)):
    return {"IMAGE_SIZE": config_get(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):
        config_set(app.state.IMAGE_SIZE, form_data.size)
        return {
            "IMAGE_SIZE": config_get(app.state.IMAGE_SIZE),
            "status": True,
        }
    else:
        raise HTTPException(
            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": config_get(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:
        config_set(app.state.IMAGE_STEPS, form_data.steps)
        return {
            "IMAGE_STEPS": config_get(app.state.IMAGE_STEPS),
            "status": True,
        }
    else:
        raise HTTPException(
            status_code=400,
            detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 50)."),
        )


@app.get("/models")
def get_models(user=Depends(get_current_user)):
    try:
        if app.state.ENGINE == "openai":
            return [
                {"id": "dall-e-2", "name": "DALL·E 2"},
                {"id": "dall-e-3", "name": "DALL·E 3"},
            ]
        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],
                )
            )

        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,
                )
            )
    except Exception as e:
        app.state.ENABLED = False
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))


@app.get("/models/default")
async def get_default_model(user=Depends(get_admin_user)):
    try:
        if app.state.ENGINE == "openai":
            return {
                "model": (
                    config_get(app.state.MODEL)
                    if config_get(app.state.MODEL)
                    else "dall-e-2"
                )
            }
        elif app.state.ENGINE == "comfyui":
            return {
                "model": (
                    config_get(app.state.MODEL) if config_get(app.state.MODEL) else ""
                )
            }
        else:
            r = requests.get(url=f"{app.state.AUTOMATIC1111_BASE_URL}/sdapi/v1/options")
            options = r.json()
            return {"model": options["sd_model_checkpoint"]}
    except Exception as e:
        config_set(app.state.ENABLED, False)
        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))


class UpdateModelForm(BaseModel):
    model: str


def set_model_handler(model: str):
    if app.state.ENGINE in ["openai", "comfyui"]:
        config_set(app.state.MODEL, model)
        return config_get(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
            )

        return options


@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
    size: Optional[str] = None
    negative_prompt: Optional[str] = None


def save_b64_image(b64_str):
    try:
        image_id = str(uuid.uuid4())

        if "," in b64_str:
            header, encoded = b64_str.split(",", 1)
            mime_type = header.split(";")[0]

            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

    except Exception as e:
        log.exception(f"Error saving image: {e}")
        return None


def save_url_image(url):
    image_id = str(uuid.uuid4())
    try:
        r = requests.get(url)
        r.raise_for_status()
        if r.headers["content-type"].split("/")[0] == "image":

            mime_type = r.headers["content-type"]
            image_format = mimetypes.guess_extension(mime_type)

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

            image_filename = f"{image_id}{image_format}"

            file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
            with open(file_path, "wb") as image_file:
                for chunk in r.iter_content(chunk_size=8192):
                    image_file.write(chunk)
            return image_filename
        else:
            log.error(f"Url does not point to an image.")
            return None

    except Exception as e:
        log.exception(f"Error saving image: {e}")
        return None


@app.post("/generations")
def generate_image(
    form_data: GenerateImageForm,
    user=Depends(get_current_user),
):

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

    r = None
    try:
        if app.state.ENGINE == "openai":

            headers = {}
            headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
            headers["Content-Type"] = "application/json"

            data = {
                "model": app.state.MODEL if app.state.MODEL != "" else "dall-e-2",
                "prompt": form_data.prompt,
                "n": form_data.n,
                "size": (
                    form_data.size
                    if form_data.size
                    else config_get(app.state.IMAGE_SIZE)
                ),
                "response_format": "b64_json",
            }

            r = requests.post(
                url=f"{app.state.OPENAI_API_BASE_URL}/images/generations",
                json=data,
                headers=headers,
            )

            r.raise_for_status()
            res = r.json()

            images = []

            for image in res["data"]:
                image_filename = save_b64_image(image["b64_json"])
                images.append({"url": f"/cache/image/generations/{image_filename}"})
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")

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

            return images

        elif app.state.ENGINE == "comfyui":

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

            if config_get(app.state.IMAGE_STEPS) is not None:
                data["steps"] = config_get(app.state.IMAGE_STEPS)

            if form_data.negative_prompt is not None:
                data["negative_prompt"] = form_data.negative_prompt

            data = ImageGenerationPayload(**data)

            res = comfyui_generate_image(
                config_get(app.state.MODEL),
                data,
                user.id,
                config_get(app.state.COMFYUI_BASE_URL),
            )
            log.debug(f"res: {res}")

            images = []

            for image in res["data"]:
                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")

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

            log.debug(f"images: {images}")
            return images
        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 config_get(app.state.IMAGE_STEPS) is not None:
                data["steps"] = config_get(app.state.IMAGE_STEPS)

            if form_data.negative_prompt is not 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()

            log.debug(f"res: {res}")

            images = []

            for image in res["images"]:
                image_filename = save_b64_image(image)
                images.append({"url": f"/cache/image/generations/{image_filename}"})
                file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")

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

            return images

    except Exception as e:
        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))