Update space

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-# 이미지 생성 함수들 import
 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
@@ -7,8 +6,6 @@ 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
-
-# 한국어 이야기 생성 함수 import
 from Text2Long_text import generate_korean_story
 
 # 모델 함수 매핑
@@ -23,25 +20,22 @@ MODEL_FUNCTIONS = {
     "한국어 긴 이야기 생성": lambda prompt: generate_korean_story(prompt, max_length=100),
 }
 
-def generate_selected_results(prompt, selected_models):
-    results = []
-    for name in selected_models:
-        gen_func = MODEL_FUNCTIONS.get(name)
-        if gen_func:
-            result = gen_func(prompt)
-            results.append((name, result))
-    # 결과별로 적절히 구분(이미지는 gallery용, 텍스트는 아래에 표시)
-    images = []
-    texts = []
-    for name, result in results:
-        if hasattr(result, "save") or hasattr(result, "show"):
-            images.append(result)  # PIL.Image로 간주
-        else:
-            texts.append(f"[{name}]\n{result}")
-    return images, "\n\n".join(texts)
+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("## 원하는 생성 모델로 이미지를 생성하거나, 긴 한국어 이야기를 받아보세요")
+    gr.Markdown("## 원하는 생성 모델로 이미지를 생성하거나, 긴 한국어 이야기를 먼저 만들고 그 이야기로 이미지를 만들 수 있습니다.")
     prompt = gr.Textbox(label="프롬프트(텍스트를 입력하세요)")
     models = gr.CheckboxGroup(
         choices=list(MODEL_FUNCTIONS.keys()),
@@ -52,7 +46,7 @@ with gr.Blocks() as demo:
     generate_btn = gr.Button("생성하기")
 
     generate_btn.click(
-        fn=generate_selected_results,
+        fn=generate_story_then_images,
         inputs=[prompt, models],
         outputs=[gallery, long_textbox]
     )