Update app.py

app.py

@@ -10,10 +10,17 @@ 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
+# ---------------------- ZeroGPU 환경에서 GPU 사용 설정 ----------------------
+# Hugging Face Spaces의 ZeroGPU 환경이면 환경 변수 ZERO_GPU가 설정되어 있다고 가정
+if os.getenv("ZERO_GPU"):
+    device = "cuda"
+    torch.cuda.set_device(0)
+else:
+    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)
 
@@ -79,7 +86,6 @@ logger = logging.getLogger("idea_generator")
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=GEMINI_API_KEY)
 
-# 슬래시("/")가 포함된 변환 문자열에서 두 옵션 중 하나만 선택하는 헬퍼 함수
 def choose_alternative(transformation):
     if "/" not in transformation:
         return transformation
@@ -101,7 +107,7 @@ def choose_alternative(transformation):
     else:
         return random.choice([left, right])
 
-# 창의적인 모델/컨셉/형상 변화 아이디어를 위한 카테고리
+# 창의적인 모델/컨셉/형상 변화 아이디어를 위한 카테고리 (총 15개)
 physical_transformation_categories = {
     "공간 이동": [
         "앞/뒤 이동", "좌/우 이동", "위/아래 이동", "세로축 회전(고개 끄덕임)", 
@@ -209,7 +215,7 @@ physical_transformation_categories = {
         "음향 반사/흡수", "음향 도플러 효과", "음파 간섭", "음향 공진",
         "진동 패턴 변화", "타악 효과", "음향 피드백", "음향 차폐/증폭",
         "소리 지향성", "음향 왜곡", "비트 생성", "하모닉스 생성", "주파수 변조",
-        "음향 충격파", "음향 필터링", "음파 전파 패턴", "진동 댐핑"
+        "음향 충격파", "음향 필터링"
     ],
     
     "생물학적 변화": [
@@ -244,7 +250,7 @@ physical_transformation_categories = {
 }
 
 ##############################################################################
-# Gemini API 호출 함수 (예: gemini-2.0-flash-thinking-exp-01-21 -> 다른 모델 사용 시 수정)
+# Gemini API 호출 함수 (모델: gemini-2.0-flash-thinking-exp-01-21)
 ##############################################################################
 def query_gemini_api(prompt):
     try:
@@ -253,15 +259,13 @@ def query_gemini_api(prompt):
         try:
             if hasattr(response, 'text'):
                 return response.text
-            
             if hasattr(response, 'candidates') and response.candidates:
-                if len(response.candidates) > 0:
-                    candidate = response.candidates[0]
-                    if hasattr(candidate, 'content'):
-                        content = candidate.content
-                        if hasattr(content, 'parts') and content.parts:
-                            if len(content.parts) > 0:
-                                return content.parts[0].text
+                candidate = response.candidates[0]
+                if hasattr(candidate, 'content'):
+                    content = candidate.content
+                    if hasattr(content, 'parts') and content.parts:
+                        if len(content.parts) > 0:
+                            return content.parts[0].text
             if hasattr(response, 'parts') and response.parts:
                 if len(response.parts) > 0:
                     return response.parts[0].text
@@ -290,48 +294,42 @@ def enhance_with_llm(base_description, obj_name, category):
     return query_gemini_api(prompt)
 
 ##############################################################################
-# 단일 키워드에 대한 "창의적 변화 아이디어" 생성
+# 선택된 카테고리만을 대상으로 한 창의적 변화 아이디어 생성 함수들
 ##############################################################################
-def generate_single_object_transformations(obj):
-    results = {}
-    for category, transformations in physical_transformation_categories.items():
-        transformation = choose_alternative(random.choice(transformations))
-        base_description = f"{obj}이(가) {transformation} 현상을 보인다"
-        results[category] = {"base": base_description, "enhanced": None}
-    return results
+def generate_single_object_transformation_for_category(obj, selected_category):
+    transformations = physical_transformation_categories.get(selected_category)
+    if not transformations:
+        return {}
+    transformation = choose_alternative(random.choice(transformations))
+    base_description = f"{obj}이(가) {transformation} 현상을 보인다"
+    return {selected_category: {"base": base_description, "enhanced": None}}
 
-##############################################################################
-# 두 키워드에 대한 "창의적 변화 아이디어" 생성
-##############################################################################
-def generate_two_objects_interaction(obj1, obj2):
-    results = {}
-    for category, transformations in physical_transformation_categories.items():
-        transformation = choose_alternative(random.choice(transformations))
-        template = random.choice([
-            "{obj1}이(가) {obj2}에 결합하여 {change}가 발생했다",
-            "{obj1}과(와) {obj2}이(가) 충돌하면서 {change}가 일어났다"
-        ])
-        base_description = template.format(obj1=obj1, obj2=obj2, change=transformation)
-        results[category] = {"base": base_description, "enhanced": None}
-    return results
+def generate_two_objects_interaction_for_category(obj1, obj2, selected_category):
+    transformations = physical_transformation_categories.get(selected_category)
+    if not transformations:
+        return {}
+    transformation = choose_alternative(random.choice(transformations))
+    template = random.choice([
+        "{obj1}이(가) {obj2}에 결합하여 {change}가 발생했다",
+        "{obj1}과(와) {obj2}이(가) 충돌하면서 {change}가 일어났다"
+    ])
+    base_description = template.format(obj1=obj1, obj2=obj2, change=transformation)
+    return {selected_category: {"base": base_description, "enhanced": None}}
 
-##############################################################################
-# 세 키워드에 대한 "창의적 변화 아이디어" 생성
-##############################################################################
-def generate_three_objects_interaction(obj1, obj2, obj3):
-    results = {}
-    for category, transformations in physical_transformation_categories.items():
-        transformation = choose_alternative(random.choice(transformations))
-        template = random.choice([
-            "{obj1}, {obj2}, {obj3}이(가) 삼각형 구조로 결합하여 {change}가 발생했다",
-            "{obj1}이(가) {obj2}와(과) {obj3} 사이에서 매개체 역할을 하며 {change}를 촉진했다"
-        ])
-        base_description = template.format(obj1=obj1, obj2=obj2, obj3=obj3, change=transformation)
-        results[category] = {"base": base_description, "enhanced": None}
-    return results
+def generate_three_objects_interaction_for_category(obj1, obj2, obj3, selected_category):
+    transformations = physical_transformation_categories.get(selected_category)
+    if not transformations:
+        return {}
+    transformation = choose_alternative(random.choice(transformations))
+    template = random.choice([
+        "{obj1}, {obj2}, {obj3}이(가) 삼각형 구조로 결합하여 {change}가 발생했다",
+        "{obj1}이(가) {obj2}와(과) {obj3} 사이에서 매개체 역할을 하며 {change}를 촉진했다"
+    ])
+    base_description = template.format(obj1=obj1, obj2=obj2, obj3=obj3, change=transformation)
+    return {selected_category: {"base": base_description, "enhanced": None}}
 
 ##############################################################################
-# 생성된 기본 설명을 LLM을 통해 확장
+# 생성된 기본 설명을 LLM을 통해 확장 (선택된 카테고리만 해당)
 ##############################################################################
 def enhance_descriptions(results, objects):
     obj_name = " 및 ".join([obj for obj in objects if obj])
@@ -340,17 +338,17 @@ def enhance_descriptions(results, objects):
     return results
 
 ##############################################################################
-# 사용자 입력(최대 3개 키워드)에 따라 창의적 변화 아이디어 생성
+# 사용자 입력(최대 3개 키워드)와 선택된 카테고리에 따라 창의적 변화 아이디어 생성
 ##############################################################################
-def generate_transformations(text1, text2=None, text3=None):
+def generate_transformations(text1, text2, text3, selected_category):
     if text2 and text3:
-        results = generate_three_objects_interaction(text1, text2, text3)
+        results = generate_three_objects_interaction_for_category(text1, text2, text3, selected_category)
         objects = [text1, text2, text3]
     elif text2:
-        results = generate_two_objects_interaction(text1, text2)
+        results = generate_two_objects_interaction_for_category(text1, text2, selected_category)
         objects = [text1, text2]
     else:
-        results = generate_single_object_transformations(text1)
+        results = generate_single_object_transformation_for_category(text1, selected_category)
         objects = [text1]
     return enhance_descriptions(results, objects)
 
@@ -378,13 +376,8 @@ def process_inputs(text1, text2, text3, selected_category, progress=gr.Progress(
     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:
-        return "선택한 카테고리가 결과에 존재하지 않습니다."
+    # 선택된 카테고리만 대상으로 아이디어 생성
+    results = generate_transformations(text1, text2, text3, selected_category)
     
     progress(0.8, desc="결과 포맷팅 중...")
     formatted = format_results(results)
@@ -395,9 +388,7 @@ def process_inputs(text1, text2, text3, selected_category, progress=gr.Progress(
 # 새로운 통합 함수: 아이디어 텍스트 생성 및 이미지 생성
 ##############################################################################
 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
 
@@ -444,7 +435,7 @@ with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 및 디자
                 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>