ui/image_gen_tab.py

"""
이미지 생성 UI 컴포넌트 및 기능
"""

import gradio as gr
import time
from config import (
    HIGH_PRICE_OPTION,
    LOW_PRICE_OPTION,
    CUSTOM,
    GEN_IMAGE_DESC,
    WORLDCUPABLE_IMAGE_CNT,
    imageStyleMap,
)
from image_generator import genPromptAndImage


def create_image_tab():
    """이미지 생성 탭 UI 생성 함수"""
    with gr.Tab("이미지 생성"):
        current_img_url = gr.State("")
        selectedImageIndex = gr.State(0)
        entryImageUrls = gr.State([])  # 저장된 URL 목록

        with gr.Column():  # Column을 사용하여 수직 배열
            with gr.Row():
                with gr.Column():  # Column을 사용하여 수직 배열
                    userPrompt = gr.TextArea(
                        label="교재 지문 (필수)",
                        value=GEN_IMAGE_DESC,
                        max_lines=5,
                    )
                    additionalComment = gr.Textbox(
                        label="추가 희망사항 (옵셔널)", lines=2, value=""
                    )
                with gr.Column():
                    with gr.Accordion("상세 설정", open=False):
                        with gr.Column():  # Column을 사용하여 수직 배열
                            with gr.Row():
                                highImgCnt = gr.Textbox(
                                    label="비싼 이미지",
                                    value="0",
                                    interactive=False,
                                )
                                lowImgCnt = gr.Textbox(
                                    label="싼 이미지", value="0", interactive=False
                                )
                                totalPrice = gr.Textbox(
                                    label="총 이미지 생성 비용(원)",
                                    value="0",
                                    interactive=False,
                                )
                            genModelKind = gr.Radio(
                                [LOW_PRICE_OPTION, HIGH_PRICE_OPTION],
                                label="이미지 생성 모델",
                                value=LOW_PRICE_OPTION,
                            )
                            with gr.Row():
                                aspectRatio = gr.Radio(
                                    ["1:1", "4:3", "16:9", CUSTOM],
                                    label="이미지 비율",
                                    value="16:9",
                                    scale=4,
                                )
                                imgWidth = gr.Number(
                                    label="너비 (32의배수)",
                                    value=1440,
                                    scale=1,
                                    visible=False,
                                    maximum=1440,
                                    minimum=256,
                                    step=32,
                                )
                                imgHeight = gr.Number(
                                    label="높이 (32의배수)",
                                    value=768,
                                    scale=1,
                                    visible=False,
                                    maximum=1440,
                                    minimum=256,
                                    step=32,
                                )

                            def custom_aspect_ratio(ratio):
                                if ratio == CUSTOM:
                                    return (
                                        gr.update(
                                            value=HIGH_PRICE_OPTION,
                                            interactive=False,
                                        ),
                                        gr.update(visible=True),
                                        gr.update(visible=True),
                                    )
                                else:
                                    return (
                                        gr.update(interactive=True),
                                        gr.update(visible=False),
                                        gr.update(visible=False),
                                    )

                            aspectRatio.change(
                                fn=custom_aspect_ratio,
                                inputs=aspectRatio,
                                outputs=[genModelKind, imgWidth, imgHeight],
                            )

                            imageFormatList = ["png", "jpg", "webp"]
                            imageFormatRadio = gr.Radio(
                                imageFormatList, label="이미지 포멧", value="webp"
                            )

                            styleList = list(imageStyleMap.keys())
                            imageStyleRadio = gr.Radio(
                                styleList, label="이미지 스타일", value="알아서"
                            )
                            with gr.Accordion("Prompt info", open=False):
                                fluxPrompt = gr.Textbox(
                                    label="Flux Prompt", value="", lines=10
                                )

            with gr.Row():
                genSmartImage = gr.Button("교재 지문 삽화 이미지 생성!")
                removeEntryWorldcupBtn = gr.Button("현재 이미지 제거")

            with gr.Column():
                progressBar = gr.Image(
                    label=None, interactive=False, visible=False, height=28
                )
                entryListGallery = gr.Gallery(
                    label="이미지 월드컵 진출 이미지",
                    preview=True,
                    allow_preview=True,
                    interactive=False,
                )
                with gr.Row():
                    entryCount = gr.Markdown("## 이미지 월드컵 진출 이미지 수: 0")
                    goImageWorldCupBtn = gr.Button(
                        "이미지 월드컵 하러가기", visible=False
                    )

            @entryListGallery.select(outputs=selectedImageIndex)
            def setSelectedImageIndex(evt: gr.SelectData, state=None):
                try:
                    if evt is None:
                        return 0
                    return evt.index
                except:
                    return 0

            def removeSelectedEntry(entryList, selectedIndex):
                if selectedIndex is None or len(entryList) == 0:
                    return gr.update(value=entryList), entryList

                del entryList[selectedIndex]
                return (
                    entryList,
                    entryList,
                    f"## 이미지 월드컵 진출 이미지 수: {len(entryList)}",
                )

            # 왠지 모르게 마지막 요소를 지우면 preview=True 가 안되서 이렇게 함.
            def reSelectedEntry(entryList):
                return gr.update(selected_index=len(entryList) - 1, preview=True)

            removeEntryWorldcupBtn.click(
                fn=removeSelectedEntry,
                inputs=[entryImageUrls, selectedImageIndex],
                outputs=[entryListGallery, entryImageUrls, entryCount],
            ).then(
                fn=reSelectedEntry,
                inputs=entryImageUrls,
                outputs=entryListGallery,
            )

            # 이벤트 처리.
            # 입력값이 변경될 때마다 Flux Prompt 초기화
            for component in [
                userPrompt,
                aspectRatio,
                imageStyleRadio,
                additionalComment,
            ]:
                component.change(fn=lambda x: "", inputs=fluxPrompt, outputs=fluxPrompt)
            genImageInputs = [
                userPrompt,
                additionalComment,
                fluxPrompt,
                aspectRatio,
                imageStyleRadio,
                imageFormatRadio,
                genModelKind,
                highImgCnt,
                lowImgCnt,
                imgWidth,
                imgHeight,
            ]

            def show_progress():
                return gr.update(label="이미지 생성 중", visible=True)

            def hide_progress():
                return gr.update(visible=False)

            genSmartImage.click(fn=show_progress, outputs=progressBar).then(
                fn=genPromptAndImage,
                inputs=genImageInputs,
                outputs=[
                    fluxPrompt,
                    highImgCnt,
                    lowImgCnt,
                    totalPrice,
                    current_img_url,
                    progressBar,
                ],
            ).then(
                fn=lambda entryList, imgUrl: (
                    gr.update(
                        value=[*entryList, imgUrl],
                        preview=True,
                        selected_index=len(entryList),
                    ),
                    entryList + [imgUrl],
                    len(entryList),
                    gr.update(visible=False),
                ),
                inputs=[entryImageUrls, current_img_url],
                outputs=[
                    entryListGallery,
                    entryImageUrls,
                    selectedImageIndex,
                    progressBar,
                ],
            ).then(
                fn=lambda entryImageUrls: f"## 이미지 월드컵 진출 이미지 수: {len(entryImageUrls)}",
                inputs=entryImageUrls,
                outputs=entryCount,
            ).then(
                fn=lambda entryCountTxt, entryList: (
                    [
                        entryCountTxt + " (이미지 월드컵 가능!)",
                        gr.update(visible=False),
                        gr.Info(
                            f"{len(entryList)}개의 이미지가 만들어졌습니다. '이미지 월드컵'을 시작하실 수 있습니다!",
                            duration=5,
                        ),
                    ]
                    if len(entryList) in WORLDCUPABLE_IMAGE_CNT
                    else [entryCountTxt, gr.update(visible=False)]
                ),
                inputs=[entryCount, entryImageUrls],
                outputs=[entryCount, goImageWorldCupBtn],
            )

    return entryImageUrls