image_generator.py

"""
이미지 생성 관련 기능을 담당하는 모듈
"""

import replicate
from config import (
    imageStyleMap,
    imageModelMap,
    HIGH_PRICE_OPTION,
    LOW_PRICE_OPTION,
    HIGH_PRICE,
    LOW_PRICE,
    CUSTOM,
)


def to_string(output):
    """
    output 을 하나의 문자열로 합치기.
    """
    combined_result = ""
    for item in output:
        combined_result += str(item)
    return combined_result


def genFluxPrompt(userPrompt, additionalComment, ratio, imageStyle):
    """
    userPrompt 와 additionalComment 를 받아서 비율에 맞는 이미지 생성을 위한 prompt 생성.
    """
    imageStyle = imageStyleMap[imageStyle]
    # print("genFluxPrompt ", userPrompt, ratio, imageStyle)
    # The ibm-granite/granite-3.1-8b-instruct model can stream output as it's running.

    additionalComment = (
        f"\n# Important note: {additionalComment}\n" if additionalComment else ""
    )

    prompt = (
        f"Please create a prompt that will generate a best quality image that can express the following sentence in English.(DO NOT INCLUDE ANY KOREAN LANGUAGE.)"
        f"\n# Image generation info\n"
        f"\n- Aspect ratio: {ratio}\n"
        f"\n- Image style: {imageStyle}\n"
        "\n# Description for the image.\n" + userPrompt + additionalComment
    )

    print("genFluxPrompt prompt:", prompt)

    result = replicate.run(
        "anthropic/claude-3.5-haiku",
        input={
            "top_k": 50,
            "top_p": 0.9,
            "prompt": prompt,
            "max_tokens": 256,
            "min_tokens": 0,
            "temperature": 0.6,
            "system_prompt": (
                "You are a professional prompt engineer specializing in creating prompts for txt2img AI model."
                "You always create prompts that can extract the impressive output."
                "Make sure to emphasize the intention for 'Important note:' section in the prompts if it exists."
                "Do not generate negative prompts!"
            ),
            "presence_penalty": 0,
            "frequency_penalty": 0,
        },
    )
    prompt = to_string(result)
    print("genFluxPrompt prompt:", prompt)
    return prompt


def genImage(
    prompt, ratio, imageModelKind, imgWidth, imgHeight, imageNum=1, imageFormat="webp"
):
    """
    prompt 를 받아서 비율에 맞는 이미지 imageNum 만큼 생성.
    """

    imageModel = imageModelMap[imageModelKind]

    if ratio == "custom":
        input = {
            "prompt": prompt,
            "aspect_ratio": ratio,  # gradio radio 버튼에서 선택된 값
            "width": int(imgWidth),
            "height": int(imgHeight),
            "aspect_ratio": "custom",
            "output_format": imageFormat,
            "output_quality": 90,
        }
        # custom 일 때는 비싼 모델로 생성.
        imageModel = imageModelMap[HIGH_PRICE_OPTION]
    else:
        input = {
            "prompt": prompt,
            "aspect_ratio": ratio,  # gradio radio 버튼에서 선택된 값
            "output_format": imageFormat,
            "output_quality": 90,
            "num_outputs": imageNum,
        }

    print("genImage input:", input)

    imageModel = imageModelMap[imageModelKind]
    result = replicate.run(imageModel, input=input)

    print("Generated Image Info:", result)
    # imgUrls = [str(img) for img in result]
    # imgUrl = str(result) if type(result) != list else str(result[0])
    imgUrl = str(result) if not isinstance(result, list) else str(result[0])
    if imageNum > 1:
        imgUrls = [str(img) for img in result]
        print("genImage imgUrls:", imgUrls)
        return imgUrls
    else:
        print("genImage imgUrl:", imgUrl)
        return imgUrl


def reduxImage(prompt, url):
    """
    url 의 이미지를 재생성.

    output = replicate.run(
    "black-forest-labs/flux-1.1-pro",
        input={
            "width": 1024,
            "height": 480,
            "prompt": "one boy and 3 girls hands and hands playing in the school.",
            "aspect_ratio": "custom",
            "image_prompt": "https://replicate.delivery/xezq/CLKVkX1QXeXrMy1tH43zJmCk6RgBszqmQdK45AOUTedCkyUUA/tmpucuderjn.webp",
            "output_format": "webp",
            "output_quality": 80,
            "safety_tolerance": 2,
            "prompt_upsampling": True
        }
    )
    """
    result = replicate.run(
        "black-forest-labs/flux-1.1-pro",
        input={
            "prompt": prompt,
            "aspect_ratio": "16:9",
            "image_prompt": url,
            "output_format": "webp",
            "output_quality": 80,
            "safety_tolerance": 2,  # 1 is most strict and 6 is most permissive
            "prompt_upsampling": True,
        },
    )
    return str(result)


def genPromptAndImage(
    userPrompt,
    additionalComment,
    fluxPrompt,
    ratio,
    imageStyle,
    imageFormat,
    imageModel,
    highImageCnt,
    lowImageCnt,
    imgWidth,
    imgHeight,
):
    """
    fluxPrompt 가 빈 문자열이면 새로운 prompt 를 생성하고,
    빈 문자열이 아니면 기존의 prompt 를 사용해서 이미지 생성.
    """
    print(
        "genPromptAndImage:",
        ratio,
        imageStyle,
        imageFormat,
        imageModel,
        highImageCnt,
        lowImageCnt,
    )

    if ratio == CUSTOM:
        ratio = "custom"

    if fluxPrompt == "":
        fluxPrompt = genFluxPrompt(userPrompt, additionalComment, ratio, imageStyle)

    imgUrl = genImage(
        fluxPrompt, ratio, imageModel, imgWidth, imgHeight, 1, imageFormat
    )

    highImageCnt = int(highImageCnt)
    lowImageCnt = int(lowImageCnt)
    if imageModel == HIGH_PRICE_OPTION:
        highImageCnt += 1
    else:
        lowImageCnt += 1

    totalPrice = highImageCnt * HIGH_PRICE + lowImageCnt * LOW_PRICE
    return fluxPrompt, highImageCnt, lowImageCnt, totalPrice, imgUrl, None