Update app.py

app.py

@@ -5,6 +5,65 @@ import time
 import logging
 import google.generativeai as genai
 
+import torch
+import numpy as np
+from diffusers import DiffusionPipeline
+from transformers import pipeline as hf_pipeline
+
+# ---------------------- 이미지 생성 관련 설정 ----------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device=="cuda" else torch.float32
+
+# 한국어-영어 번역 모델 로드 (장치에 따라 CPU 또는 GPU 사용)
+translator = hf_pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device=0 if device=="cuda" else -1)
+
+# Diffusion Pipeline 로드 (예: FLUX.1-schnell 모델)
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=dtype
+).to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+def contains_korean(text):
+    for char in text:
+        if ord('가') <= ord(char) <= ord('힣'):
+            return True
+    return False
+
+def generate_design_image(prompt, seed=42, randomize_seed=True, width=1024, height=1024, num_inference_steps=4):
+    """
+    생성된 확장 아이디어 텍스트(prompt)를 입력받아,
+    필요시 한국어를 영어로 번역한 후 DiffusionPipeline으로 이미지를 생성합니다.
+    """
+    original_prompt = prompt
+    translated = False
+    
+    # 한국어가 포함되어 있으면 영어로 번역
+    if contains_korean(prompt):
+        translation = translator(prompt)
+        prompt = translation[0]['translation_text']
+        translated = True
+
+    # 랜덤 시드 설정
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    image = pipe(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=0.0
+    ).images[0]
+    
+    return image
+
+# ---------------------- Gemini API 및 아이디어 생성 관련 기존 코드 ----------------------
 # 로깅 설정
 logging.basicConfig(
     level=logging.INFO,
@@ -189,12 +248,8 @@ physical_transformation_categories = {
 ##############################################################################
 def query_gemini_api(prompt):
     try:
-        # 예시: 기존 gemini-2.0... 대신, 다른 모델이 필요하다면 교체하세요.
         model = genai.GenerativeModel('gemini-2.0-flash-thinking-exp-01-21')
-        
         response = model.generate_content(prompt)
-
-        # 응답 구조 방어적으로 처리
         try:
             if hasattr(response, 'text'):
                 return response.text
@@ -210,13 +265,10 @@ def query_gemini_api(prompt):
             if hasattr(response, 'parts') and response.parts:
                 if len(response.parts) > 0:
                     return response.parts[0].text
-
             return "Unable to generate a response. API response structure is different than expected."
-            
         except Exception as inner_e:
             logger.error(f"Error processing response: {inner_e}")
             return f"An error occurred while processing the response: {str(inner_e)}"
-            
     except Exception as e:
         logger.error(f"Error calling Gemini API: {e}")
         if "API key not valid" in str(e):
@@ -224,7 +276,7 @@ def query_gemini_api(prompt):
         return f"An error occurred while calling the API: {str(e)}"
 
 ##############################################################################
-# 설명 확장 함수: "모델/컨셉/형상의 변화에 대한 이해와 혁신 포인트, 기능성 등을 중심"으로
+# 설명 확장 함수
 ##############################################################################
 def enhance_with_llm(base_description, obj_name, category):
     prompt = f"""
@@ -238,7 +290,7 @@ def enhance_with_llm(base_description, obj_name, category):
     return query_gemini_api(prompt)
 
 ##############################################################################
-# 단일 키워드(오브젝트)에 대한 "창의적 변화 아이디어" 생성
+# 단일 키워드에 대한 "창의적 변화 아이디어" 생성
 ##############################################################################
 def generate_single_object_transformations(obj):
     results = {}
@@ -283,10 +335,8 @@ def generate_three_objects_interaction(obj1, obj2, obj3):
 ##############################################################################
 def enhance_descriptions(results, objects):
     obj_name = " 및 ".join([obj for obj in objects if obj])
-    
     for category, result in results.items():
         result["enhanced"] = enhance_with_llm(result["base"], obj_name, category)
-    
     return results
 
 ##############################################################################
@@ -302,7 +352,6 @@ def generate_transformations(text1, text2=None, text3=None):
     else:
         results = generate_single_object_transformations(text1)
         objects = [text1]
-    
     return enhance_descriptions(results, objects)
 
 ##############################################################################
@@ -315,7 +364,7 @@ def format_results(results):
     return formatted
 
 ##############################################################################
-# Gradio UI에서 호출할 함수
+# Gradio UI에서 호출할 함수 (아이디어 텍스트 생성)
 ##############################################################################
 def process_inputs(text1, text2, text3, selected_category, progress=gr.Progress()):
     text1 = text1.strip() if text1 else None
@@ -325,20 +374,13 @@ def process_inputs(text1, text2, text3, selected_category, progress=gr.Progress(
     if not text1:
         return "오류: 최소 하나의 키워드를 입력해주세요."
     
-    keyword_info = f"키워드: {text1}"
-    if text2:
-        keyword_info += f", {text2}"
-    if text3:
-        keyword_info += f", {text3}"
-    
     progress(0.05, desc="아이디어 생성 준비 중...")
-    time.sleep(0.3)  # 시각적 효과를 위한 짧은 지연
-    
-    progress(0.1, desc="창의적인 모델/컨셉/형상 변화 아이디어 생성 시작...")
+    time.sleep(0.3)
+    progress(0.1, desc="창의적인 아이디어 생성 시작...")
     
     results = generate_transformations(text1, text2, text3)
     
-    # 선택한 카테고리에 해당하는 결과만 필터링
+    # 선택한 카테고리 결과만 필터링
     if selected_category in results:
         results = {selected_category: results[selected_category]}
     else:
@@ -346,10 +388,19 @@ def process_inputs(text1, text2, text3, selected_category, progress=gr.Progress(
     
     progress(0.8, desc="결과 포맷팅 중...")
     formatted = format_results(results)
-    
     progress(1.0, desc="완료!")
     return formatted
 
+##############################################################################
+# 새로운 통합 함수: 아이디어 텍스트 생성 및 이미지 생성
+##############################################################################
+def process_all(text1, text2, text3, selected_category, progress=gr.Progress()):
+    # 확장 아이디어 텍스트 생성
+    idea_result = process_inputs(text1, text2, text3, selected_category, progress)
+    # 생성된 아이디어를 그대로 이미지 생성 프롬프트로 사용
+    image_result = generate_design_image(idea_result, seed=42, randomize_seed=True, width=1024, height=1024, num_inference_steps=4)
+    return idea_result, image_result
+
 ##############################################################################
 # API 키 경고 메시지
 ##############################################################################
@@ -361,7 +412,7 @@ def get_warning_message():
 ##############################################################################
 # Gradio UI
 ##############################################################################
-with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 생성기",
+with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 및 디자인 생성기",
                theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
     
     gr.HTML("""
@@ -377,27 +428,23 @@ with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 생성기",
     </style>
     """)
     
-    gr.Markdown("# 🚀 키워드 기반 창의적 변화 아이디어 생성기")
-    gr.Markdown("입력한 **키워드**(최대 3개)와 **카테고리**를 바탕으로, **창의적인 모델/컨셉/형상 변화**에 대한 이해와 **혁신 포인트**, **기능성** 등을 중심으로 확장된 아이디어를 제시합니다.")
+    gr.Markdown("# 🚀 키워드 기반 창의적 변화 아이디어 및 디자인 생성기")
+    gr.Markdown("입력한 **키워드**(최대 3개)와 **카테고리**를 바탕으로, 창의적인 모델/컨셉/형상 변화 아이디어를 생성하고, 해당 확장 아이디어를 프롬프트로 하여 디자인 이미지를 생성합니다.")
     
     warning = gr.Markdown(get_warning_message())
     
-    # 좌측 입력 영역
     with gr.Row():
         with gr.Column(scale=1):
             text_input1 = gr.Textbox(label="키워드 1 (필수)", placeholder="예: 스마트폰")
             text_input2 = gr.Textbox(label="키워드 2 (선택)", placeholder="예: 인공지능")
             text_input3 = gr.Textbox(label="키워드 3 (선택)", placeholder="예: 헬스케어")
-            # 카테고리 선택 드롭다운 추가
             category_dropdown = gr.Dropdown(
                 label="카테고리 선택",
                 choices=list(physical_transformation_categories.keys()),
                 value=list(physical_transformation_categories.keys())[0],
                 info="출력할 카테고리를 선택하세요."
             )
-            
-            status_msg = gr.Markdown("💡 '아이디어 생성하기' 버튼을 클릭하면 아이디어 생성이 시작됩니다.")
-            
+            status_msg = gr.Markdown("���� '아이디어 생성하기' 버튼을 클릭하면 아이디어 생성과 함께 디자인 이미지가 생성됩니다.")
             processing_indicator = gr.HTML("""
             <div style="display: flex; justify-content: center; align-items: center; margin: 10px 0;">
                 <div style="border: 5px solid #f3f3f3; border-top: 5px solid #3498db; border-radius: 50%; width: 30px; height: 30px; animation: spin 2s linear infinite;"></div>
@@ -410,13 +457,12 @@ with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 생성기",
             }
             </style>
             """, visible=False)
-            
             submit_button = gr.Button("아이디어 생성하기", variant="primary")
         
-        # 우측 출력 영역
         with gr.Column(scale=2):
             idea_output = gr.Markdown(label="아이디어 결과")
-
+            generated_image = gr.Image(label="생성된 디자인 이미지", type="pil")
+    
     gr.Examples(
         examples=[
             ["스마트폰", "", "", list(physical_transformation_categories.keys())[0]],
@@ -439,9 +485,9 @@ with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 생성기",
         inputs=None,
         outputs=processing_indicator
     ).then(
-        fn=process_inputs,
+        fn=process_all,
         inputs=[text_input1, text_input2, text_input3, category_dropdown],
-        outputs=idea_output
+        outputs=[idea_output, generated_image]
     ).then(
         fn=hide_processing_indicator,
         inputs=None,