Update app.py

app.py

@@ -12,7 +12,7 @@ import logging
 from google import genai
 from google.genai import types
 
-# .env 파일에 저장된 환경변수 로드 (pip install python-dotenv 필요)
+# .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -29,11 +29,22 @@ def save_binary_file(file_name, data):
     logger.debug(f"파일 저장 완료: {file_name}")
 
 
-def generate(text, composite_file, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
-    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', composite_file: '{composite_file}', "
-                 f"background_file: '{background_file}', style_file: '{style_file}', 모델: '{model}'")
+def generate(text, files_list, model="gemini-2.0-flash-exp-image-generation"):
+    """
+    여러 이미지와 텍스트 프롬프트를 사용하여 Gemini API로 이미지를 생성하는 함수
+    
+    Args:
+        text (str): 이미지 생성에 사용할 텍스트 프롬프트
+        files_list (list): 이미지 파일 경로 리스트
+        model (str): 사용할 Gemini 모델명
+    
+    Returns:
+        str: 생성된 이미지의 파일 경로 또는 None (오류 발생 시)
+    """
+    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일 수: {len(files_list)}, 모델: '{model}'")
+
     try:
-        # 환경변수에서 API 키 불러오기
+        # API 키는 환경변수에서 불러옴
         effective_api_key = os.environ.get("GEMINI_API_KEY")
         if effective_api_key:
             logger.debug("환경변수에서 API 키 불러옴")
@@ -44,32 +55,40 @@ def generate(text, composite_file, background_file=None, style_file=None, model=
         client = genai.Client(api_key=effective_api_key)
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
-        # 업로드할 파일들 리스트 생성
-        files = []
-        # 원본 이미지 업로드 (필수)
-        composite_upload = client.files.upload(file=composite_file)
-        files.append(composite_upload)
-        logger.debug(f"원본 이미지 업로드 완료. URI: {composite_upload.uri}, MIME 타입: {composite_upload.mime_type}")
-
-        # 배경 이미지 업로드 (선택)
-        if background_file is not None:
-            background_upload = client.files.upload(file=background_file)
-            files.append(background_upload)
-            logger.debug(f"배경 이미지 업로드 완료. URI: {background_upload.uri}, MIME 타입: {background_upload.mime_type}")
-        # 스타일 이미지 업로드 (선택)
-        if style_file is not None:
-            style_upload = client.files.upload(file=style_file)
-            files.append(style_upload)
-            logger.debug(f"스타일 이미지 업로드 완료. URI: {style_upload.uri}, MIME 타입: {style_upload.mime_type}")
-
-        # 파일 업로드 결과를 바탕으로 콘텐츠 파트 생성
+        # 여러 파일 업로드
+        uploaded_files = []
+        for file_path in files_list:
+            if file_path and os.path.exists(file_path):
+                uploaded_file = client.files.upload(file=file_path)
+                uploaded_files.append(uploaded_file)
+                logger.debug(f"파일 업로드 완료. URI: {uploaded_file.uri}, MIME 타입: {uploaded_file.mime_type}")
+            else:
+                logger.warning(f"파일이 존재하지 않거나 None입니다: {file_path}")
+        
+        if not uploaded_files:
+            logger.error("업로드할 파일이 없습니다.")
+            return None
+
+        # 파트 리스트 생성 (모든 이미지 + 텍스트 프롬프트)
         parts = []
-        for file_obj in files:
-            parts.append(types.Part.from_uri(file_uri=file_obj.uri, mime_type=file_obj.mime_type))
-        # 마지막에 텍스트 프롬프트 추가
+        for uploaded_file in uploaded_files:
+            parts.append(
+                types.Part.from_uri(
+                    file_uri=uploaded_file.uri,
+                    mime_type=uploaded_file.mime_type,
+                )
+            )
+        
+        # 텍스트 프롬프트 추가
         parts.append(types.Part.from_text(text=text))
-
-        contents = [types.Content(role="user", parts=parts)]
+        
+        # 컨텐츠 객체 생성
+        contents = [
+            types.Content(
+                role="user",
+                parts=parts,
+            ),
+        ]
         logger.debug(f"컨텐츠 객체 생성 완료: {contents}")
 
         generate_content_config = types.GenerateContentConfig(
@@ -77,7 +96,10 @@ def generate(text, composite_file, background_file=None, style_file=None, model=
             top_p=0.95,
             top_k=40,
             max_output_tokens=8192,
-            response_modalities=["image", "text"],
+            response_modalities=[
+                "image",
+                "text",
+            ],
             response_mime_type="text/plain",
         )
         logger.debug(f"생성 설정: {generate_content_config}")
@@ -108,7 +130,8 @@ def generate(text, composite_file, background_file=None, style_file=None, model=
 
                 logger.debug(f"Raw chunk: {chunk}")
 
-        del files
+        # 업로드된 파일 정보 정리
+        del uploaded_files
         logger.debug("업로드된 파일 정보 삭제 완료.")
         return temp_path
 
@@ -117,42 +140,84 @@ def generate(text, composite_file, background_file=None, style_file=None, model=
         return None  # 오류 발생 시 None 반환
 
 
-def process_image_and_prompt(composite_pil, background_pil, style_pil, prompt):
-    logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
-    try:
-        # 원본 이미지 저장
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            composite_path = tmp.name
-            composite_pil.save(composite_path)
-            logger.debug(f"원본 이미지 저장 완료: {composite_path}")
-
-        background_path = None
-        style_path = None
-
-        # 배경 이미지 저장 (입력된 경우)
-        if background_pil is not None:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp2:
-                background_path = tmp2.name
-                background_pil.save(background_path)
-                logger.debug(f"배경 이미지 저장 완료: {background_path}")
-
-        # 스타일 이미지 저장 (입력된 경우)
-        if style_pil is not None:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp3:
-                style_path = tmp3.name
-                style_pil.save(style_path)
-                logger.debug(f"스타일 이미지 저장 완료: {style_path}")
+def save_image_temp(image):
+    """이미지를 임시 파일로 저장하는 유틸리티 함수"""
+    if image is None:
+        return None
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+        temp_path = tmp.name
+        if isinstance(image, Image.Image):
+            if image.mode == "RGBA":
+                # 알파 채널이 있는 경우 처리
+                image = image.convert("RGB")
+            image.save(temp_path)
+        else:
+            logger.warning(f"지원되지 않는 이미지 타입: {type(image)}")
+            return None
+        
+        logger.debug(f"이미지 저장 완료: {temp_path}")
+        return temp_path
 
+
+def process_with_background_and_style(main_image, background_image, style_image, prompt):
+    """
+    주 이미지, 배경 이미지, 스타일 이미지를 프롬프트와 함께 처리하는 함수
+    
+    Args:
+        main_image (PIL.Image): 주 대상 이미지
+        background_image (PIL.Image): 배경 이미지
+        style_image (PIL.Image): 스타일 참조 이미지
+        prompt (str): 이미지 편집 프롬프트
+    
+    Returns:
+        list: 처리된 이미지 목록
+    """
+    logger.debug(f"process_with_background_and_style 함수 시작 - 프롬프트: '{prompt}'")
+    try:
+        # 각 이미지를 임시 파일로 저장
+        files_list = []
+        
+        # 주 이미지 저장
+        main_path = save_image_temp(main_image)
+        if main_path:
+            files_list.append(main_path)
+        
+        # 배경 이미지 저장 (있는 경우)
+        if background_image is not None:
+            bg_path = save_image_temp(background_image)
+            if bg_path:
+                files_list.append(bg_path)
+        
+        # 스타일 이미지 저장 (있는 경우)
+        if style_image is not None:
+            style_path = save_image_temp(style_image)
+            if style_path:
+                files_list.append(style_path)
+        
+        # 이미지가 하나도 없는 경우 처리
+        if not files_list:
+            logger.error("처리할 이미지가 없습니다.")
+            return []
+        
+        # 프롬프트 보강 (배경 및 스타일 관련 정보 추가)
+        enhanced_prompt = prompt
+        if background_image is not None and style_image is not None:
+            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지는 배경으로 사용해주세요. 세 번째 이미지의 스타일을 적용해주세요."
+        elif background_image is not None:
+            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지는 배경으로 사용해주세요."
+        elif style_image is not None:
+            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지의 스타일을 적용해주세요."
+        
+        logger.debug(f"보강된 프롬프트: {enhanced_prompt}")
         model = "gemini-2.0-flash-exp-image-generation"
-        gemma_edited_image_path = generate(text=prompt,
-                                           composite_file=composite_path,
-                                           background_file=background_path,
-                                           style_file=style_path,
-                                           model=model)
-
-        if gemma_edited_image_path:
-            logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
-            result_img = Image.open(gemma_edited_image_path)
+
+        # Gemini API 호출
+        result_path = generate(text=enhanced_prompt, files_list=files_list, model=model)
+
+        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")
             return [result_img]
@@ -161,7 +226,7 @@ def process_image_and_prompt(composite_pil, background_pil, style_pil, prompt):
             return []  # 오류 시 빈 리스트 반환
 
     except Exception as e:
-        logger.exception("process_image_and_prompt 함수에서 오류 발생:")
+        logger.exception("process_with_background_and_style 함수에서 오류 발생:")
         return []  # 오류 시 빈 리스트 반환
 
 
@@ -175,45 +240,78 @@ with gr.Blocks() as demo:
             </div>
             <div>
                 <h1>Gemini를 이용한 이미지 편집</h1>
+                <p>배경 이미지와 스타일 이미지를 추가로 적용할 수 있습니다.</p>
                 <p>Gemini API 키는 환경변수(GEMINI_API_KEY)로 설정되어 있습니다.</p>
             </div>
         </div>
         """
     )
-    gr.Markdown("원본 이미지, 배경 이미지, 스타일 이미지를 업로드하고, 편집할 내용을 입력하세요.")
+    gr.Markdown("이미지를 업로드하고, 편집할 내용을 입력하세요. 배경 이미지와 스타일 이미지는 선택사항입니다.")
 
     with gr.Row():
         with gr.Column():
-            composite_input = gr.Image(type="pil", label="원본 이미지", image_mode="RGBA")
-            background_input = gr.Image(type="pil", label="배경 이미지", image_mode="RGBA")
-            style_input = gr.Image(type="pil", label="스타일 이미지", image_mode="RGBA")
+            # 입력 컴포넌트
+            main_image_input = gr.Image(type="pil", label="주 이미지 업로드", image_mode="RGB")
+            background_image_input = gr.Image(type="pil", label="배경 이미지 업로드 (선택사항)", image_mode="RGB")
+            style_image_input = gr.Image(type="pil", label="스타일 이미지 업로드 (선택사항)", image_mode="RGB")
+            
             prompt_input = gr.Textbox(
-                lines=2,
+                lines=3,
                 placeholder="편집할 내용을 입력하세요...",
                 label="편집 프롬프트"
             )
+            
+            # 프롬프트 예시
+            gr.Markdown("### 프롬프트 예시")
+            gr.Markdown("- '주 이미지를 배경 이미지 위에 자연스럽게 배치해주세요'")
+            gr.Markdown("- '주 이미지를 스타일 이미지의 화풍으로 변환해주세요'")
+            gr.Markdown("- '주 이미지의 배경을 두 번째 이미지로 대체하고, 세 번째 이미지의 색감을 적용해주세요'")
+            
             submit_btn = gr.Button("이미지 편집 실행")
+            
         with gr.Column():
+            # 출력 컴포넌트
             output_gallery = gr.Gallery(label="편집 결과")
+            output_info = gr.Textbox(label="처리 결과 정보", lines=2)
 
+    # 처리 함수 연결
     submit_btn.click(
-        fn=process_image_and_prompt,
-        inputs=[composite_input, background_input, style_input, prompt_input],
-        outputs=output_gallery,
+        fn=process_with_background_and_style,
+        inputs=[main_image_input, background_image_input, style_image_input, prompt_input],
+        outputs=[output_gallery],
+    ).then(
+        lambda: "편집이 완료되었습니다. 결과를 확인해주세요.",
+        None,
+        output_info
     )
 
-# --- 테스트 코드 ---
-dummy_composite = Image.new("RGBA", (100, 100), color="red")
-dummy_background = Image.new("RGBA", (100, 100), color="green")
-dummy_style = Image.new("RGBA", (100, 100), color="blue")
-dummy_prompt = "이미지에 새로운 스타일을 적용해줘"
-
-logger.info("process_image_and_prompt 함수를 직접 호출합니다...")
-result = process_image_and_prompt(dummy_composite, dummy_background, dummy_style, dummy_prompt)
+    # 예제 이미지 추가
+    gr.Examples(
+        examples=[
+            ["example_main.jpg", "example_background.jpg", "example_style.jpg", "주 이미지의 인물을 배경 이미지에 합성하고, 스타일 이미지의 화풍을 적용해주세요"],
+            ["example_main.jpg", None, "example_style.jpg", "이미지를 스타일 이미지의 화풍으로 변환해주세요"],
+            ["example_main.jpg", "example_background.jpg", None, "주 이미지의 배경을 제거하고 두 번째 이미지 위에 합성해주세요"],
+        ],
+        inputs=[main_image_input, background_image_input, style_image_input, prompt_input],
+    )
 
-if result:
-    logger.info(f"직접 호출 성공. 결과: {result}")
-else:
-    logger.error("직접 호출 실패.")
 
-demo.launch(share=True)
+# 애플리케이션 실행
+if __name__ == "__main__":
+    # 실제 운영에서는 테스트 코드 제거
+    # 테스트용 더미 이미지 (실제 이미지로 대체 가능)
+    # dummy_image = Image.new("RGB", (100, 100), color="red")
+    # dummy_bg = Image.new("RGB", (100, 100), color="blue")
+    # dummy_style = Image.new("RGB", (100, 100), color="green")
+    # dummy_prompt = "이미지를 합성해줘"
+    
+    # logger.info("process_with_background_and_style 함수를 직접 호출합니다...")
+    # result = process_with_background_and_style(dummy_image, dummy_bg, dummy_style, dummy_prompt)
+    
+    # if result:
+    #     logger.info(f"직접 호출 성공. 결과: {result}")
+    # else:
+    #     logger.error("직접 호출 실패.")
+    
+    # Gradio 앱 실행
+    demo.launch(share=True)