Update app.py

app.py

@@ -9,15 +9,12 @@ import base64
 import mimetypes
 import logging
 
-import google.generativeai as genai
-from google.generativeai import types
+from google import genai
+from google.genai import types
 
 # .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
-try:
-    from dotenv import load_dotenv
-    load_dotenv()
-except ImportError:
-    logger.warning("python-dotenv가 설치되지 않았습니다. 환경변수를 직접 설정해야 합니다.")
+from dotenv import load_dotenv
+load_dotenv()
 
 # 로깅 설정 (로그 레벨: DEBUG)
 logging.basicConfig(level=logging.DEBUG,
@@ -32,19 +29,8 @@ def save_binary_file(file_name, data):
     logger.debug(f"파일 저장 완료: {file_name}")
 
 
-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}'")
+def generate(text, file_name, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
+    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일명: '{file_name}', 모델: '{model}'")
 
     try:
         # API 키는 환경변수에서 불러옴
@@ -58,34 +44,50 @@ def generate(text, files_list, model="gemini-2.0-flash-exp-image-generation"):
         client = genai.Client(api_key=effective_api_key)
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
-        # 여러 파일 업로드
+        # 원본, 배경, 스타일 이미지를 각각 업로드 (존재하는 경우)
         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
-
-        # 파트 리스트 생성 (모든 이미지 + 텍스트 프롬프트)
+        uploaded_files.append(client.files.upload(file=file_name))
+        logger.debug(f"원본 파일 업로드 완료. URI: {uploaded_files[0].uri}, MIME 타입: {uploaded_files[0].mime_type}")
+
+        if background_file is not None:
+            bg_file = client.files.upload(file=background_file)
+            uploaded_files.append(bg_file)
+            logger.debug(f"배경 파일 업로드 완료. URI: {bg_file.uri}, MIME 타입: {bg_file.mime_type}")
+        if style_file is not None:
+            style_uploaded = client.files.upload(file=style_file)
+            uploaded_files.append(style_uploaded)
+            logger.debug(f"스타일 파일 업로드 완료. URI: {style_uploaded.uri}, MIME 타입: {style_uploaded.mime_type}")
+
+        # 컨텐츠 객체 생성: 업로드한 이미지들을 순서대로(원본, 배경, 스타일) 추가하고 마지막에 텍스트 프롬프트 추가
         parts = []
-        for uploaded_file in uploaded_files:
+        # 원본 이미지 파트
+        parts.append(
+            types.Part.from_uri(
+                file_uri=uploaded_files[0].uri,
+                mime_type=uploaded_files[0].mime_type,
+            )
+        )
+        # 배경 이미지 파트 (존재 시)
+        if background_file is not None:
             parts.append(
                 types.Part.from_uri(
-                    file_uri=uploaded_file.uri,
-                    mime_type=uploaded_file.mime_type,
+                    file_uri=uploaded_files[1].uri,
+                    mime_type=uploaded_files[1].mime_type,
                 )
             )
-        
-        # 텍스트 프롬프트 추가
+        # 스타일 이미지 파트 (존재 시)
+        if style_file is not None:
+            # 배경이 없는 경우 uploaded_files[1]가 스타일 이미지가 됨에 주의
+            style_index = 2 if background_file is not None else 1
+            parts.append(
+                types.Part.from_uri(
+                    file_uri=uploaded_files[style_index].uri,
+                    mime_type=uploaded_files[style_index].mime_type,
+                )
+            )
+        # 마지막으로 텍스트 파트 추가
         parts.append(types.Part.from_text(text=text))
-        
-        # 컨텐츠 객체 생성
+
         contents = [
             types.Content(
                 role="user",
@@ -133,7 +135,6 @@ def generate(text, files_list, model="gemini-2.0-flash-exp-image-generation"):
 
                 logger.debug(f"Raw chunk: {chunk}")
 
-        # 업로드된 파일 정보 정리
         del uploaded_files
         logger.debug("업로드된 파일 정보 삭제 완료.")
         return temp_path
@@ -143,84 +144,43 @@ def generate(text, files_list, model="gemini-2.0-flash-exp-image-generation"):
         return None  # 오류 발생 시 None 반환
 
 
-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}'")
+def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pil=None):
+    logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{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}")
+        # 원본 이미지 저장
+        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}")
+
+        input_text = prompt
         model = "gemini-2.0-flash-exp-image-generation"
 
-        # Gemini API 호출
-        result_path = generate(text=enhanced_prompt, files_list=files_list, model=model)
+        gemma_edited_image_path = generate(text=input_text,
+                                           file_name=original_path,
+                                           background_file=background_path,
+                                           style_file=style_path,
+                                           model=model)
 
-        if result_path:
-            logger.debug(f"이미지 생성 완료. 경로: {result_path}")
-            result_img = Image.open(result_path)
+        if gemma_edited_image_path:
+            logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
+            result_img = Image.open(gemma_edited_image_path)
             if result_img.mode == "RGBA":
                 result_img = result_img.convert("RGB")
             return [result_img]
@@ -229,7 +189,7 @@ def process_with_background_and_style(main_image, background_image, style_image,
             return []  # 오류 시 빈 리스트 반환
 
     except Exception as e:
-        logger.exception("process_with_background_and_style 함수에서 오류 발생:")
+        logger.exception("process_image_and_prompt 함수에서 오류 발생:")
         return []  # 오류 시 빈 리스트 반환
 
 
@@ -243,78 +203,46 @@ 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():
-            # 입력 컴포넌트
-            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")
-            
+            original_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")
             prompt_input = gr.Textbox(
-                lines=3,
+                lines=2,
                 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_with_background_and_style,
-        inputs=[main_image_input, background_image_input, style_image_input, prompt_input],
-        outputs=[output_gallery],
-    ).then(
-        lambda: "편집이 완료되었습니다. 결과를 확인해주세요.",
-        None,
-        output_info
+        fn=process_image_and_prompt,
+        inputs=[original_input, prompt_input, background_input, style_input],
+        outputs=output_gallery,
     )
 
-    # 예제 이미지 추가
-    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],
-    )
+# --- 테스트 코드 ---
+# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
+dummy_original = 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_original, dummy_prompt, dummy_background, dummy_style)
 
+if result:
+    logger.info(f"직접 호출 성공. 결과: {result}")
+else:
+    logger.error("직접 호출 실패.")
 
-# 애플리케이션 실행
-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)
+demo.launch(share=True)