import gradio as gr
from model.animagine_xl import generate_animagine_xl
from model.animesai import generate_animesai
from model.generate_sdxl_with_refiner import generate_sdxl_with_refiner
from model.ghibli import generate_ghibli
from model.realistic import generate_realistic
from model.sd_turbo import generate_sd_turbo
from model.waifu import generate_waifu
from Text2Long_text import generate_korean_story

# 모델 함수 매핑

model_id = "hakurei/waifu-diffusion"

MODEL_FUNCTIONS = {
    "Animagine XL": generate_animagine_xl,
    "Animesai": generate_animesai,
    "SDXL+Refiner": generate_sdxl_with_refiner,
    "Ghibli": generate_ghibli,
    "Realistic": generate_realistic,
    "SD Turbo": generate_sd_turbo,
    "Waifu Diffusion": generate_waifu,
    "한국어 긴 이야기 생성": lambda prompt: generate_korean_story(prompt, max_length=100),
}

def generate_story_then_images(prompt, selected_models):
    # 1. 선택지 중 '한국어 긴 이야기 생성'이 포함되어 있으면 먼저 스토리를 생성
    story = ""
    if "한국어 긴 이야기 생성" in selected_models:
        story = generate_korean_story(prompt, max_length=100)
    # 2. 이미지 생성 모델들에는 이야기(스토리)가 있으면 그걸 프롬프트로 넣어줌
    img_prompt = story if story else prompt
    images = [
        MODEL_FUNCTIONS[name](img_prompt)
        for name in selected_models
        if name != "한국어 긴 이야기 생성"
    ]
    return images, story

with gr.Blocks() as demo:
    gr.Markdown("## 원하는 생성 모델로 이미지를 생성하거나, 긴 한국어 이야기를 먼저 만들고 그 이야기로 이미지를 만들 수 있습니다.")
    prompt = gr.Textbox(label="프롬프트(텍스트를 입력하세요)")
    models = gr.CheckboxGroup(
        choices=list(MODEL_FUNCTIONS.keys()),
        label="사용할 모델을 선택하세요"
    )
    gallery = gr.Gallery(label="생성 이미지 (선택한 모델별)")
    long_textbox = gr.Textbox(label="생성된 이야기(텍스트)", lines=10, interactive=False)
    generate_btn = gr.Button("생성하기")

    generate_btn.click(
        fn=generate_story_then_images,
        inputs=[prompt, models],
        outputs=[gallery, long_textbox]
    )

demo.launch()