Update app.py

app.py

@@ -3,6 +3,7 @@ import tempfile
 from PIL import Image
 import gradio as gr
 import logging
+import re
 
 from google import genai
 from google.genai import types
@@ -11,7 +12,7 @@ from google.genai import types
 from dotenv import load_dotenv
 load_dotenv()
 
-# 간단한 로깅 설정
+# 로깅 설정
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
@@ -19,6 +20,14 @@ def save_binary_file(file_name, data):
     with open(file_name, "wb") as f:
         f.write(data)
 
+def preprocess_prompt(prompt):
+    """
+    프롬프트에 있는 #1, #2, #3 참조를 적절한 텍스트로 변환
+    """
+    # #1, #2, #3 참조를 텍스트로 변환
+    processed_prompt = prompt.replace("#1", "첫 번째 이미지").replace("#2", "두 번째 이미지").replace("#3", "세 번째 이미지")
+    return processed_prompt
+
 def process_images_with_prompt(image1, image2, image3, prompt):
     """
     3개의 이미지와 프롬프트를 처리하는 함수
@@ -32,9 +41,12 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         # Gemini 클라이언트 초기화
         client = genai.Client(api_key=api_key)
         
-        # 프롬프트 기본값 설정
+        # 프롬프트 처리
         if not prompt or not prompt.strip():
             prompt = "이 이미지들을 활용하여 새로운 이미지를 생성해주세요."
+        else:
+            # #1, #2, #3 참조 처리
+            prompt = preprocess_prompt(prompt)
         
         # 컨텐츠 리스트 생성 (이미지와 프롬프트 결합)
         parts = []
@@ -43,7 +55,8 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         parts.append(types.Part.from_text(text=prompt))
         
         # 사용 가능한 이미지 추가
-        for img in [image1, image2, image3]:
+        images_added = 0
+        for idx, img in enumerate([image1, image2, image3], 1):
             if img is not None:
                 # PIL 이미지를 바이트로 변환
                 with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
@@ -54,6 +67,11 @@ def process_images_with_prompt(image1, image2, image3, prompt):
                 
                 # 이미지를 파트로 추가
                 parts.append(types.Part.from_data(data=image_bytes, mime_type="image/png"))
+                images_added += 1
+                logger.info(f"이미지 #{idx} 추가됨")
+        
+        if images_added == 0:
+            return None, "적어도 하나의 이미지를 업로드해주세요."
         
         # 생성 설정
         generate_content_config = types.GenerateContentConfig(
@@ -66,6 +84,7 @@ def process_images_with_prompt(image1, image2, image3, prompt):
             temp_path = tmp.name
             
             # Gemini 모델로 요청 전송
+            logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
             response = client.models.generate_content(
                 model="gemini-2.0-flash-exp-image-generation",
                 contents=[types.Content(role="user", parts=parts)],
@@ -73,9 +92,15 @@ def process_images_with_prompt(image1, image2, image3, prompt):
             )
             
             # 응답에서 이미지 추출
+            image_found = False
             for part in response.candidates[0].content.parts:
                 if 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, "API에서 이미지를 생성하지 못했습니다. 다른 프롬프트로 시도해보세요."
                     
             # 결과 이미지 반환
             result_img = Image.open(temp_path)
@@ -88,39 +113,72 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         logger.exception("이미지 생성 중 오류 발생:")
         return None, f"오류 발생: {str(e)}"
 
-# 간소화된 Gradio 인터페이스
+# Gradio 인터페이스
 with gr.Blocks() as demo:
-    gr.HTML("<h1>간단한 이미지 생성기</h1><p>이미지 3개와 프롬프트를 입력하세요</p>")
+    gr.HTML(
+        """
+        <div style="text-align: center; margin-bottom: 1rem;">
+            <h1>간단한 이미지 생성기</h1>
+            <p>이미지와 프롬프트를 입력하세요. #1, #2, #3으로 각 이미지를 참조할 수 있습니다.</p>
+        </div>
+        """
+    )
     
     with gr.Row():
         with gr.Column():
             # 3개의 이미지 입력
-            image1_input = gr.Image(type="pil", label="이미지 1", image_mode="RGB")
-            image2_input = gr.Image(type="pil", label="이미지 2", image_mode="RGB")
-            image3_input = gr.Image(type="pil", label="이미지 3", image_mode="RGB")
+            with gr.Row():
+                image1_input = gr.Image(type="pil", label="#1", image_mode="RGB")
+                image2_input = gr.Image(type="pil", label="#2", image_mode="RGB")
+                image3_input = gr.Image(type="pil", label="#3", image_mode="RGB")
             
             # 프롬프트 입력
             prompt_input = gr.Textbox(
                 lines=3,
-                placeholder="이 이미지들을 어떻게 변환할지 설명해주세요",
+                placeholder="프롬프트를 입력하세요. 예: '#1과 #2를 합성해 주세요' 또는 '#1의 인물을 #3의 배경에 넣어주세요'",
                 label="프롬프트"
             )
             
+            # 예시 프롬프트 버튼들
+            with gr.Row():
+                prompt1_btn = gr.Button("예시: #1 + #2 합성")
+                prompt2_btn = gr.Button("예시: #1의 인물 + #2의 배경")
+                prompt3_btn = gr.Button("예시: #3 스타일로 #1 변환")
+            
             # 생성 버튼
-            submit_btn = gr.Button("이미지 생성")
+            submit_btn = gr.Button("이미지 생성", variant="primary")
         
         with gr.Column():
             # 결과 출력
             output_image = gr.Image(label="생성된 이미지")
             output_text = gr.Textbox(label="상태 메시지")
     
-    # 버튼 클릭 이벤트
+    # 예시 프롬프트 버튼 클릭 이벤트
+    prompt1_btn.click(lambda: "#1과 #2를 자연스럽게 합성해 주세요", outputs=prompt_input)
+    prompt2_btn.click(lambda: "#1의 인물을 #2의 배경에 넣어주세요", outputs=prompt_input)
+    prompt3_btn.click(lambda: "#3의 스타일로 #1을 변환해 주세요", outputs=prompt_input)
+    
+    # 이미지 생성 버튼 클릭 이벤트
     submit_btn.click(
         fn=process_images_with_prompt,
         inputs=[image1_input, image2_input, image3_input, prompt_input],
         outputs=[output_image, output_text],
     )
 
+    gr.HTML(
+        """
+        <div style="margin-top: 1rem; padding: 1rem; background-color: #f8f9fa; border-radius: 0.5rem;">
+            <h3>사용 방법:</h3>
+            <ul>
+                <li>이미지를 #1, #2, #3 슬롯에 업로드하세요</li>
+                <li>프롬프트에서 #1, #2, #3으로 각 이미지를 참조할 수 있습니다</li>
+                <li>예시 버튼을 클릭하면 자주 사용하는 프롬프트 템플릿을 적용할 수 있습니다</li>
+                <li>모든 이미지를 업로드할 필요는 없으며, 필요한 이미지만 업로드하세요</li>
+            </ul>
+        </div>
+        """
+    )
+
 # 애플리케이션 실행
 if __name__ == "__main__":
     demo.launch(share=True)