Update app.py

app.py

@@ -1,12 +1,9 @@
 import json
 import os
 import time
-import uuid
 import tempfile
 from PIL import Image
 import gradio as gr
-import base64
-import mimetypes
 import logging
 from io import BytesIO
 
@@ -30,8 +27,8 @@ def save_binary_file(file_name, data):
     logger.debug(f"파일 저장 완료: {file_name}")
 
 
-def generate(text, original_image_path, background_image_path=None, style_image_path=None, model="gemini-2.0-flash-exp-image-generation"):
-    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 원본 파일명: '{original_image_path}', 모델: '{model}'")
+def merge_product_and_person(person_img_path, product_img_path, prompt, model="gemini-2.0-flash-exp-image-generation"):
+    logger.debug(f"merge_product_and_person 함수 시작 - 프롬프트: '{prompt}'")
     
     try:
         # API 키는 환경변수에서 불러옴
@@ -46,33 +43,12 @@ def generate(text, original_image_path, background_image_path=None, style_image_
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
         # PIL 이미지 객체로 변환
-        original_img = Image.open(original_image_path)
+        person_img = Image.open(person_img_path)
+        product_img = Image.open(product_img_path)
         
         # 컨텐츠 리스트 생성 (공식 문서 방식대로)
-        contents = []
-        
-        # 텍스트 프롬프트를 명확하게 작성
-        prompt = text
-        if background_image_path and "배경" not in text.lower():
-            prompt += " 원본 이미지의 배경을 두 번째 업로드된 이미지로 완전히 교체해 주세요. 이미지를 업데이트하고 결과를 보여주세요."
-        if style_image_path and "스타일" not in text.lower():
-            prompt += " 세 번째 이미지의 스타일을 전체적으로 적용해 주세요."
-            
-        contents.append(prompt)
-        contents.append(original_img)
-        
-        # 배경 이미지 추가 (있는 경우)
-        if background_image_path:
-            background_img = Image.open(background_image_path)
-            contents.append(background_img)
-            logger.debug("배경 이미지 추가됨")
-            
-        # 스타일 이미지 추가 (있는 경우)
-        if style_image_path:
-            style_img = Image.open(style_image_path)
-            contents.append(style_img)
-            logger.debug("스타일 이미지 추가됨")
-            
+        # 순서: 사람 이미지, 상품 이미지, 프롬프트 텍스트
+        contents = [person_img, product_img, prompt]
         logger.debug(f"컨텐츠 객체 생성 완료: {len(contents)} 아이템")
 
         # 생성 설정
@@ -81,7 +57,7 @@ def generate(text, original_image_path, background_image_path=None, style_image_
             top_p=0.95,
             top_k=40,
             max_output_tokens=8192,
-            response_modalities=["text", "image"],  # 문서에 맞게 대소문자 수정
+            response_modalities=["text", "image"],
         )
         logger.debug(f"생성 설정: {generate_content_config}")
 
@@ -89,7 +65,7 @@ def generate(text, original_image_path, background_image_path=None, style_image_
             temp_path = tmp.name
             logger.debug(f"임시 파일 생성됨: {temp_path}")
 
-            # 스트리밍 대신 단일 요청으로 변경 (이미지 생성에 더 적합)
+            # 단일 요청으로 이미지 생성
             response = client.models.generate_content(
                 model=model,
                 contents=contents,
@@ -98,12 +74,14 @@ def generate(text, original_image_path, background_image_path=None, style_image_
             
             logger.debug("응답 처리 시작...")
             
-            # 응답에서 이미지 추출
+            # 응답에서 이미지와 텍스트 추출
             image_saved = False
+            response_text = ""
+            
             for part in response.candidates[0].content.parts:
                 if hasattr(part, 'text') and part.text:
+                    response_text += part.text
                     logger.info(f"수신된 텍스트: {part.text}")
-                    print(part.text)
                 elif hasattr(part, 'inline_data') and part.inline_data:
                     save_binary_file(temp_path, part.inline_data.data)
                     logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
@@ -111,84 +89,67 @@ def generate(text, original_image_path, background_image_path=None, style_image_
 
             if not image_saved:
                 logger.warning("이미지가 생성되지 않았습니다.")
-                return None
+                return None, response_text
 
         logger.debug("이미지 생성 완료.")
-        return temp_path
+        return temp_path, response_text
 
     except Exception as e:
         logger.exception("이미지 생성 중 오류 발생:")
-        return None  # 오류 발생 시 None 반환
+        return None, str(e)  # 오류 발생 시 None과 오류 메시지 반환
 
 
-def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pil=None):
-    logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
+def process_images_and_prompt(person_pil, product_pil, prompt):
+    logger.debug(f"process_images_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
     try:
-        # 원본 이미지 저장
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            original_path = tmp.name
-            original_pil.save(original_path)
-            logger.debug(f"원본 이미지 저장 완료: {original_path}")
-
-        # 배경 이미지 저장 (제공된 경우)
-        background_path = None
-        if background_pil is not None:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_bg:
-                background_path = tmp_bg.name
-                background_pil.save(background_path)
-                logger.debug(f"배경 이미지 저장 완료: {background_path}")
-
-        # 스타일 이미지 저장 (제공된 경우)
-        style_path = None
-        if style_pil is not None:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_style:
-                style_path = tmp_style.name
-                style_pil.save(style_path)
-                logger.debug(f"스타일 이미지 저장 완료: {style_path}")
-
-        # 프롬프트 보강 (영어와 한국어 두 언어로 제공)
-        if prompt and not prompt.strip():
-            if background_path and style_path:
-                prompt = "원본 이미지의 인물을 유지하면서 배경을 두 번째 이미지로 교체하고 세 번째 이미지의 스타일을 적용해 주세요. Please replace the background with the second image while keeping the person, and apply the style of the third image."
-            elif background_path:
-                prompt = "원본 이미지의 인물을 유지하면서 배경을 두 번째 이미지로 교체해 주세요. Please replace the background with the second image while keeping the person."
-            elif style_path:
-                prompt = "원본 이미지에 두 번째 이미지의 스타일을 적용해 주세요. Please apply the style of the second image to the original image."
-
-        model = "gemini-2.0-flash-exp-image-generation"
-
-        gemma_edited_image_path = generate(
-            text=prompt,
-            original_image_path=original_path,
-            background_image_path=background_path,
-            style_image_path=style_path,
-            model=model
+        # 기본 프롬프트 설정 (비어있는 경우)
+        if not prompt or not prompt.strip():
+            prompt = "이 사람이 이 상품을 사용하는 모습을 자연스럽게 보여주세요. 상품을 잘 보이게 해주세요. Make a natural composite image showing this person using this product. Make sure the product is clearly visible."
+        
+        # 프롬프트에 영어가 없으면 영어 프롬프트 추가 (더 나은 결과를 위해)
+        if not any(ord(c) < 128 for c in prompt):
+            prompt += " Create a realistic composite image of this person with this product."
+        
+        # 이미지 저장
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_person:
+            person_path = tmp_person.name
+            person_pil.save(person_path)
+            logger.debug(f"사람 이미지 저장 완료: {person_path}")
+            
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_product:
+            product_path = tmp_product.name
+            product_pil.save(product_path)
+            logger.debug(f"상품 이미지 저장 완료: {product_path}")
+
+        # 이미지 합성 실행
+        result_path, response_text = merge_product_and_person(
+            person_img_path=person_path,
+            product_img_path=product_path,
+            prompt=prompt
         )
 
-        if gemma_edited_image_path:
-            logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
-            result_img = Image.open(gemma_edited_image_path)
+        # 이미지 반환 및 임시 파일 정리
+        if result_path:
+            logger.debug(f"이미지 생성 완료. 경로: {result_path}")
+            result_img = Image.open(result_path)
             if result_img.mode == "RGBA":
                 result_img = result_img.convert("RGB")
             
             # 임시 파일 정리
             try:
-                os.unlink(original_path)
-                if background_path:
-                    os.unlink(background_path)
-                if style_path:
-                    os.unlink(style_path)
+                os.unlink(person_path)
+                os.unlink(product_path)
             except Exception as e:
                 logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
             
-            return [result_img]
+            return [result_img], response_text
         else:
-            logger.error("generate 함수에서 None 반환됨.")
-            return []  # 오류 시 빈 리스트 반환
+            logger.error("merge_product_and_person 함수에서 None 반환됨.")
+            return [], response_text  # 오류 시 빈 리스트 반환
 
     except Exception as e:
-        logger.exception("process_image_and_prompt 함수에서 오류 발생:")
-        return []  # 오류 시 빈 리스트 반환
+        logger.exception("process_images_and_prompt 함수에서 오류 발생:")
+        return [], str(e)  # 오류 시 빈 리스트와 오류 메시지 반환
 
 
 # --- Gradio 인터페이스 구성 ---
@@ -200,46 +161,47 @@ with gr.Blocks() as demo:
                 <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
             </div>
             <div>
-                <h1>Gemini를 이용한 이미지 편집</h1>
+                <h1>Gemini를 이용한 상품-사람 이미지 합성</h1>
                 <p>Gemini API 키는 환경변수(GEMINI_API_KEY)로 설정되어 있습니다.</p>
             </div>
         </div>
         """
     )
-    gr.Markdown("원본, 배경, 스타일 이미지를 업로드하고, 편집할 내용을 입력하세요.")
+    gr.Markdown("사람 이미지와 상품 이미지를 업로드하고, 어떻게 합성할지 설명해주세요.")
 
     with gr.Row():
         with gr.Column():
-            original_input = gr.Image(type="pil", label="원본 이미지", image_mode="RGB")
-            background_input = gr.Image(type="pil", label="배경 이미지", image_mode="RGB")
-            style_input = gr.Image(type="pil", label="스타일 이미지", image_mode="RGB")
+            person_input = gr.Image(type="pil", label="사람 이미지", image_mode="RGB")
+            product_input = gr.Image(type="pil", label="상품 이미지", image_mode="RGB")
             prompt_input = gr.Textbox(
                 lines=2,
-                placeholder="편집할 내용을 입력하세요...",
-                label="편집 프롬프트"
+                placeholder="합성 방법을 설명해주세요. (예: '이 사람이 상품을 들고 있는 모습' 또는 '이 사람이 이 상품을 사용하는 모습')",
+                label="합성 방법 설명"
             )
-            submit_btn = gr.Button("이미지 편집 실행")
+            submit_btn = gr.Button("이미지 합성 실행")
         with gr.Column():
-            output_gallery = gr.Gallery(label="편집 결과")
-            output_text = gr.Textbox(label="API 응답 텍스트", visible=True)
+            output_gallery = gr.Gallery(label="합성 결과")
+            output_text = gr.Textbox(label="AI 응답 텍스트", visible=True)
 
     submit_btn.click(
-        fn=process_image_and_prompt,
-        inputs=[original_input, prompt_input, background_input, style_input],
-        outputs=output_gallery,
+        fn=process_images_and_prompt,
+        inputs=[person_input, product_input, prompt_input],
+        outputs=[output_gallery, output_text],
     )
 
     gr.HTML("""
     <div style="margin-top: 20px; padding: 10px; background-color: #f8f9fa; border-radius: 8px;">
         <h3>사용 팁:</h3>
         <ul>
-            <li><strong>프롬프트 작성 예시:</strong> "원본 인물을 유지하면서 배경을 두 번째 이미지로 교체해 주세요."</li>
-            <li><strong>배경 교체 명확화:</strong> "배경만 완전히 교체하고 원본 인물은 그대로 유지해 주세요."</li>
-            <li><strong>스타일 적용:</strong> "세 번째 이미지의 예술 스타일을 전체 이미지에 적용해 주세요."</li>
+            <li><strong>자연스러운 합성:</strong> "이 사람이 이 상품을 사용하는 자연스러운 모습을 만들어주세요."</li>
+            <li><strong>특정 위치 지정:</strong> "이 사람이 이 상품을 손에 들고 있는 모습을 보여주세요."</li>
+            <li><strong>상품 강조:</strong> "이 사람이 이 상품을 사용하는 모습을 보여주되, 상품이 잘 보이도록 해주세요."</li>
             <li><strong>영어 프롬프트:</strong> 더 나은 결과를 위해 영어와 한국어를 함께 사용해 보세요.</li>
+            <li><strong>구체적인 지시:</strong> 합성 방법을 구체적으로 설명할수록 더 좋은 결과를 얻을 수 있습니다.</li>
         </ul>
     </div>
     """)
 
 # --- 실행 ---
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True)