Update app.py

app.py

@@ -4,9 +4,9 @@ from PIL import Image
 import os
 import json
 import tempfile
+import logging
 import re
 import time
-import logging
 
 # 로깅 설정
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -157,7 +157,7 @@ def generate_prompt_with_gemini(product_name, background_info, additional_info="
         )
         response = model.generate_content(
             prompt_request,
-            generation_config=genai.GenerationConfig(
+            generation_config=genai.types.GenerationConfig(
                 temperature=0.7,
                 top_p=0.95,
                 top_k=64,
@@ -171,7 +171,8 @@ def generate_prompt_with_gemini(product_name, background_info, additional_info="
     except Exception as e:
         return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"
 
-# 이미지 생성에 필요한 함수들
+# 이미지 생성 관련 새로운 함수들
+
 def save_binary_file(file_name, data):
     with open(file_name, "wb") as f:
         f.write(data)
@@ -185,14 +186,15 @@ def translate_prompt_to_english(prompt):
     prompt = prompt.replace("#3", "IMAGE_TAG_THREE")
     
     try:
-        if not GEMINI_API_KEY:
+        api_key = GEMINI_API_KEY
+        if not api_key:
             logger.error("Gemini API 키가 설정되지 않았습니다.")
             prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
             prompt = prompt.replace("IMAGE_TAG_TWO", "#2")
             prompt = prompt.replace("IMAGE_TAG_THREE", "#3")
             return prompt
             
-        model = genai.GenerativeModel('gemini-2.0-flash')
+        client = genai.Client(api_key=api_key)
         translation_prompt = f"""
         Translate the following Korean text to English:
         
@@ -203,9 +205,10 @@ def translate_prompt_to_english(prompt):
         """
         
         logger.info(f"Translation prompt: {translation_prompt}")
-        response = model.generate_content(
-            translation_prompt,
-            generation_config=genai.GenerationConfig(
+        response = client.generate_content(
+            model="gemini-2.0-flash",
+            contents=[translation_prompt],
+            generation_config=genai.types.GenerationConfig(
                 temperature=0.2,
                 top_p=0.95,
                 top_k=40,
@@ -213,7 +216,10 @@ def translate_prompt_to_english(prompt):
             )
         )
         
-        translated_text = response.text
+        translated_text = ""
+        for part in response.text:
+            if hasattr(part, 'text') and part.text:
+                translated_text += part.text
         
         if translated_text.strip():
             translated_text = translated_text.replace("IMAGE_TAG_ONE", "#1")
@@ -284,10 +290,11 @@ def preprocess_prompt(prompt, image1, image2, image3):
 
 def generate_with_images(prompt, images, variation_index=0):
     try:
-        if not GEMINI_API_KEY:
+        api_key = GEMINI_API_KEY
+        if not api_key:
             return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
 
-        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
+        client = genai.Client(api_key=api_key)
         logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")
 
         variation_suffixes = [
@@ -308,9 +315,11 @@ def generate_with_images(prompt, images, variation_index=0):
                 contents.append(img)
                 logger.info(f"이미지 #{idx} 추가됨")
 
-        response = model.generate_content(
+        response = client.generate_content(
+            model="gemini-2.0-flash-exp-image-generation",
             contents=contents,
-            generation_config=genai.GenerationConfig(
+            generation_config=genai.types.GenerationConfig(
+                response_modalities=['Text', 'Image'],
                 temperature=1,
                 top_p=0.95,
                 top_k=40,
@@ -322,26 +331,19 @@ def generate_with_images(prompt, images, variation_index=0):
             temp_path = tmp.name
             result_text = ""
             image_found = False
-            
-            if hasattr(response, 'candidates') and response.candidates:
-                candidate = response.candidates[0]
-                if hasattr(candidate, 'content') and candidate.content:
-                    for part in candidate.content.parts:
-                        if hasattr(part, 'text') and part.text:
-                            result_text += part.text
-                            logger.info(f"응답 텍스트: {part.text}")
-                        elif hasattr(part, 'inline_data') and part.inline_data:
-                            save_binary_file(temp_path, part.inline_data.data)
-                            image_found = True
-                            logger.info("응답에서 이미지 추출 성공")
-            
+            for part in response.candidates[0].content.parts:
+                if hasattr(part, 'text') and part.text:
+                    result_text += part.text
+                    logger.info(f"응답 텍스트: {part.text}")
+                elif hasattr(part, 'inline_data') and part.inline_data:
+                    save_binary_file(temp_path, part.inline_data.data)
+                    image_found = True
+                    logger.info("응답에서 이미지 추출 성공")
             if not image_found:
                 return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
-            
             result_img = Image.open(temp_path)
             if result_img.mode == "RGBA":
                 result_img = result_img.convert("RGB")
-            
             return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
     except Exception as e:
         logger.exception("이미지 생성 중 오류 발생:")
@@ -427,20 +429,26 @@ def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progres
     return results[0], results[1], results[2], results[3], combined_status, combined_prompts
 
 def create_app():
-    with gr.Blocks() as demo:
-        gr.Markdown("# 고급 상품 이미지 배경 프롬프트 생성기")
-        gr.Markdown("상품 이미지를 업로드하고 옵션을 선택하면 고품질 상업용 프롬프트가 생성됩니다.")
+    with gr.Blocks(title="이미지 생성 애플리케이션") as demo:
+        gr.Markdown("# 상품 이미지 배경 프롬프트 생성기 & 이미지 생성기")
+        gr.Markdown("상품 이미지를 업로드하고 옵션을 선택하면 고품질 상업용 프롬프트가 생성되고 이미지가 생성됩니다.")
         
+        # 상품명 섹션
         with gr.Row():
-            with gr.Column():
-                image_input = gr.Image(type="pil", label="상품 이미지 업로드")
-                
-                gr.Markdown("## 배경 유형 선택")
-                background_type = gr.Radio(
-                    choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "추상/특수 배경"],
-                    label="배경 유형",
-                    value="심플 배경"
-                )
+            product_name = gr.Textbox(label="상품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 텀블러 등", interactive=True)
+        
+        # 배경 유형을 선택하기 위한 라디오 버튼
+        background_type = gr.Radio(
+            choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "추상/특수 배경"],
+            label="배경 유형",
+            value="심플 배경"
+        )
+        
+        # 각 배경 유형에 맞는 드롭다운 컴포넌트들 (처음에는 심플 배경만 표시)
+        with gr.Row():
+            # 좌측 컬럼 - 입력 컨트롤
+            with gr.Column(scale=1):
+                image_input = gr.Image(label="상품 이미지 업로드", type="pil")
                 
                 # 심플 배경 선택 드롭다운
                 simple_dropdown = gr.Dropdown(
@@ -469,7 +477,7 @@ def create_app():
                     interactive=True
                 )
                 
-                # 실내 환경 선택 드롭다운
+# 실내 환경 선택 드롭다운
                 indoor_dropdown = gr.Dropdown(
                     choices=list(INDOOR_BACKGROUNDS.keys()),
                     value=list(INDOOR_BACKGROUNDS.keys())[0] if INDOOR_BACKGROUNDS else None,
@@ -487,39 +495,40 @@ def create_app():
                     interactive=True
                 )
                 
-                # 상품명 입력
-                product_name = gr.Textbox(label="상품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 텀블러 등")
+                # 이미지 입력 컴포넌트 (참조 이미지용)
+                gr.Markdown("## 추가 참조 이미지 (옵션)")
+                image2_input = gr.Image(label="참조 이미지 #2 (선택사항)", type="pil", image_mode="RGB")
+                image3_input = gr.Image(label="참조 이미지 #3 (선택사항)", type="pil", image_mode="RGB")
                 
                 # 추가 요청사항
                 additional_info = gr.Textbox(
                     label="추가 요청사항 (선택사항)",
                     placeholder="예: 고급스러운 느낌, 밝은 조명, 자연스러운 보조적인 객체를 추가해주세요 등",
-                    lines=3
+                    lines=3,
+                    interactive=True
                 )
                 
-                generate_prompt_btn = gr.Button("프롬프트 생성")
+                # 버튼 영역
+                with gr.Row():
+                    generate_prompt_btn = gr.Button("프롬프트 생성")
+                    submit_single_btn = gr.Button("이미지 생성 (1장)")
+                    submit_btn = gr.Button("이미지 생성 (4장)")
             
-            with gr.Column():
-                prompt_output = gr.Textbox(label="생성된 프롬프트", lines=6)
+            # 우측 컬럼 - 출�� 결과
+            with gr.Column(scale=1):
+                prompt_output = gr.Textbox(label="생성된 프롬프트", lines=10)
                 
-                gr.Markdown("## 이미지 생성")
+                # 이미지 출력 영역
+                gr.Markdown("## 생성된 이미지")
                 with gr.Row():
-                    generate_single_btn = gr.Button("이미지 생성 (1장)")
-                    generate_multiple_btn = gr.Button("이미지 생성 (4장)")
-        
-        gr.Markdown("## 생성된 이미지")
-        with gr.Row():
-            with gr.Column():
-                output_image1 = gr.Image(label="이미지 #1")
-            with gr.Column():
-                output_image2 = gr.Image(label="이미지 #2")
-        with gr.Row():
-            with gr.Column():
-                output_image3 = gr.Image(label="이미지 #3")
-            with gr.Column():
-                output_image4 = gr.Image(label="이미지 #4")
-        
-        status_output = gr.Textbox(label="상태 메시지", lines=2)
+                    output_image1 = gr.Image(label="이미지 #1", type="pil")
+                    output_image2 = gr.Image(label="이미지 #2", type="pil")
+                with gr.Row():
+                    output_image3 = gr.Image(label="이미지 #3", type="pil")
+                    output_image4 = gr.Image(label="이미지 #4", type="pil")
+                    
+                output_text = gr.Textbox(label="상태 메시지", lines=2)
+                prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", visible=True, lines=2)
         
         # 배경 유형에 따라 드롭다운 표시 업데이트
         def update_dropdowns(bg_type):
@@ -558,7 +567,7 @@ def create_app():
                     "name": nature,
                     "english": NATURE_BACKGROUNDS.get(nature, "natural environment")
                 }
-elif bg_type == "실내 환경":
+            elif bg_type == "실내 환경":
                 return {
                     "category": "실내 환경",
                     "name": indoor,
@@ -578,12 +587,9 @@ elif bg_type == "실내 환경":
                 }
         
         # 프롬프트 생성 함수
-        def generate_prompt(image, bg_type, simple, studio, nature, indoor, abstract, product_text, additional_text):
-            if image is None:
-                gr.Warning("이미지를 업로드해주세요.")
-                return "이미지를 업로드해주세요."
-            
-            product_text = product_text.strip() or "제품"
+        def generate_prompt(bg_type, simple, studio, nature, indoor, abstract, product_text, additional_text):
+            if not product_text.strip():
+                product_text = "제품"
             
             # 배경 정보 가져오기
             background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, abstract)
@@ -600,7 +606,6 @@ API 키 설정 방법:
 2. 코드 내 직접 입력: GEMINI_API_KEY = "your-api-key"
 키 발급: https://makersuite.google.com/
 """
-                    return prompt
                 
                 return prompt
             except Exception as e:
@@ -608,38 +613,10 @@ API 키 설정 방법:
                 gr.Error(error_msg)
                 return error_msg
         
-        # 단일 이미지 생성 함수
-        def generate_single_image(image, prompt):
-            if image is None:
-                return None, None, None, None, "이미지를 업로드해주세요."
-            
-            if not prompt or prompt.strip() == "":
-                return None, None, None, None, "프롬프트를 생성해주세요."
-            
-            if not GEMINI_API_KEY:
-                return None, None, None, None, "Gemini API 키가 설정되지 않았습니다. 환경 변수를 확인해주세요."
-            
-            try:
-                # 영어 프롬프트로 변환
-                if re.search("[가-힣]", prompt):
-                    prompt = translate_prompt_to_english(prompt)
-                
-                # 이미지 1개 생성 (variation_index=0)
-                result_img, status, _ = process_images_with_prompt(image, None, None, prompt, 0, 3)
-                
-                if result_img is None:
-                    return None, None, None, None, status
-                
-                return result_img, None, None, None, status
-            except Exception as e:
-                error_msg = f"이미지 생성 중 오류 발생: {str(e)}"
-                return None, None, None, None, error_msg
-        
-        # 프롬프트 생성 버튼 이벤트 연결
+        # 프롬프트 생성 버튼 클릭 이벤트
         generate_prompt_btn.click(
             fn=generate_prompt,
             inputs=[
-                image_input,
                 background_type,
                 simple_dropdown,
                 studio_dropdown,
@@ -652,18 +629,18 @@ API 키 설정 방법:
             outputs=prompt_output
         )
         
-        # 단일 이미지 생성 버튼 이벤트 연결
-        generate_single_btn.click(
-            fn=generate_single_image,
-            inputs=[image_input, prompt_output],
-            outputs=[output_image1, output_image2, output_image3, output_image4, status_output]
+        # 단일 이미지 생성 버튼 클릭 이벤트
+        submit_single_btn.click(
+            fn=lambda image1, image2, image3, prompt: process_images_with_prompt(image1, image2, image3, prompt, 0),
+            inputs=[image_input, image2_input, image3_input, prompt_output],
+            outputs=[output_image1, output_text, prompt_display],
         )
         
-        # 여러 이미지 생성 버튼 이벤트 연결
-        generate_multiple_btn.click(
-            fn=lambda img, prompt: generate_multiple_images(img, None, None, prompt),
-            inputs=[image_input, prompt_output],
-            outputs=[output_image1, output_image2, output_image3, output_image4, status_output, prompt_output]
+        # 4장 이미지 생성 버튼 클릭 이벤트
+        submit_btn.click(
+            fn=generate_multiple_images,
+            inputs=[image_input, image2_input, image3_input, prompt_output],
+            outputs=[output_image1, output_image2, output_image3, output_image4, output_text, prompt_display],
         )
     
     return demo