GLASSES_AI_SERVER/rest/app/routes/services.py

# -*- coding: utf-8 -*-
"""
@File: services.py
@Date: 2020-09-14
@author: A2TEC
@section MODIFYINFO 수정정보
- 수정자/수정일   : 수정내역
- 2022-01-14/hsj100@a2tec.co.kr : refactoring
@brief: services api
"""

import requests, json, traceback
from fastapi import APIRouter, Depends, Body
from starlette.requests import Request
from typing import Annotated, List

from rest.app.common import consts
from rest.app import models as M
from rest.app.utils.date_utils import D
from custom_logger.custom_log import custom_logger as LOG

from custom_apps.dalle3.custom_dalle import DallEArgument,dalle3_generate_image
from custom_apps.imagen.custom_imagen import imagen_generate_image
from rest.app.utils.parsing_utils import download_range


router = APIRouter(prefix="/services")


@router.post("/imageGenerate/dalle3", summary="이미지 생성(AI) - DALL-E 3", response_model=M.ImageGenerateRes)
async def dalle3(request: Request, request_body_info: M.ImageGenerateReq):
    """
    ## 이미지 생성(AI) - DALL-E 3
    > DALL-E 3 AI를 이용하여 이미지 생성
    
    ### Requriements
    > - 쿠키 정보 설정(https://github.com/acheong08/BingImageCreator) 추후 set api 추가 예정 -> 현재 고정값
    > - const.py 에 지정한 OUTPUT_FOLDER 하위에 dalle 폴더가 있어야함.
    
    ## 정보
    > 오류 발생시 오류 발생한 파일은 에러 메세지에만 남기고 저장은 안함
    
    """
    response = M.ImageGenerateRes()
    try:
        if not download_range(request_body_info.downloadCount):
            raise Exception(f"downloadCount is 1~4 (current value = {request_body_info.downloadCount})")


        args = DallEArgument(
            prompt=request_body_info.prompt,
            download_count=request_body_info.downloadCount
            )       
        
        info = dalle3_generate_image(args)
        
        if info.get_error_messages():
            error_message = f"파일생성 error: {info.get_error_messages()}"
            LOG.error(error_message)
            return response.set_error(error=error_message, img_len=info.get_counter())
        
        return response.set_message(img_len=info.get_counter())
       
    except Exception as e:
        LOG.error(traceback.format_exc())
        return response.set_error(e)
    

@router.post("/imageGenerate/imagen", summary="이미지 생성(AI) - imagen", response_model=M.ImageGenerateRes)
async def imagen(request: Request, request_body_info: M.ImageGenerateReq):
    """
    ## 이미지 생성(AI) - imagen
    > imagen AI를 이용하여 이미지 생성
    
    ### Requriements
    > - googlecli 설치(https://cloud.google.com/sdk/docs/install?hl=ko#linux)
    > - const.py 에 지정한 OUTPUT_FOLDER 하위에 imagen 폴더가 있어야함.
    
    """
    response = M.ImageGenerateRes()
    try:
        
        if not download_range(request_body_info.downloadCount):
            raise Exception(f"downloadCount is 1~4 (current value = {request_body_info.downloadCount})")
        
        img_length = imagen_generate_image(prompt=request_body_info.prompt,
                              download_count=request_body_info.downloadCount
                              )
        
        return response.set_message(img_len=img_length)
    
    except Exception as e:
        LOG.error(traceback.format_exc())
        return response.set_error(error=e)