Update app.py

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 import base64
 import mimetypes
 import logging
+from io import BytesIO
 
 from google import genai
 from google.genai import types
@@ -29,9 +30,9 @@ def save_binary_file(file_name, data):
     logger.debug(f"파일 저장 완료: {file_name}")
 
 
-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}'")
-
+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}'")
+    
     try:
         # API 키는 환경변수에서 불러옴
         effective_api_key = os.environ.get("GEMINI_API_KEY")
@@ -44,68 +45,43 @@ def generate(text, file_name, background_file=None, style_file=None, model="gemi
         client = genai.Client(api_key=effective_api_key)
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
-        # 원본, 배경, 스타일 이미지를 각각 업로드 (존재하는 경우)
-        uploaded_files = []
-        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 = []
-        # 원본 이미지 파트
-        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_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",
-                parts=parts,
-            ),
-        ]
-        logger.debug(f"컨텐츠 객체 생성 완료: {contents}")
-
+        # PIL 이미지 객체로 변환
+        original_img = Image.open(original_image_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("스타일 이미지 추가됨")
+            
+        logger.debug(f"컨텐츠 객체 생성 완료: {len(contents)} 아이템")
+
+        # 생성 설정
         generate_content_config = types.GenerateContentConfig(
             temperature=1,
             top_p=0.95,
             top_k=40,
             max_output_tokens=8192,
-            response_modalities=[
-                "image",
-                "text",
-            ],
-            response_mime_type="text/plain",
+            response_modalities=["text", "image"],  # 문서에 맞게 대소문자 수정
         )
         logger.debug(f"생성 설정: {generate_content_config}")
 
@@ -113,30 +89,31 @@ def generate(text, file_name, background_file=None, style_file=None, model="gemi
             temp_path = tmp.name
             logger.debug(f"임시 파일 생성됨: {temp_path}")
 
-            response_stream = client.models.generate_content_stream(
+            # 스트리밍 대신 단일 요청으로 변경 (이미지 생성에 더 적합)
+            response = client.models.generate_content(
                 model=model,
                 contents=contents,
                 config=generate_content_config,
             )
-
-            logger.debug("응답 스트림 처리 시작...")
-            for chunk in response_stream:
-                if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
-                    logger.warning("chunk에 후보, 컨텐츠, 또는 파트가 없습니다. 건너뜁니다.")
-                    continue
-
-                inline_data = chunk.candidates[0].content.parts[0].inline_data
-                if inline_data:
-                    save_binary_file(temp_path, inline_data.data)
-                    logger.info(f"MIME 타입 {inline_data.mime_type}의 파일이 저장됨: {temp_path} (프롬프트: {text})")
-                else:
-                    logger.info(f"수신된 텍스트: {chunk.text}")
-                    print(chunk.text)
-
-                logger.debug(f"Raw chunk: {chunk}")
-
-        del uploaded_files
-        logger.debug("업로드된 파일 정보 삭제 완료.")
+            
+            logger.debug("응답 처리 시작...")
+            
+            # 응답에서 이미지 추출
+            image_saved = False
+            for part in response.candidates[0].content.parts:
+                if hasattr(part, 'text') and 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}")
+                    image_saved = True
+
+            if not image_saved:
+                logger.warning("이미지가 생성되지 않았습니다.")
+                return None
+
+        logger.debug("이미지 생성 완료.")
         return temp_path
 
     except Exception as e:
@@ -169,20 +146,41 @@ def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pi
                 style_pil.save(style_path)
                 logger.debug(f"스타일 이미지 저장 완료: {style_path}")
 
-        input_text = prompt
+        # 프롬프트 보강 (영어와 한국어 두 언어로 제공)
+        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=input_text,
-                                           file_name=original_path,
-                                           background_file=background_path,
-                                           style_file=style_path,
-                                           model=model)
+        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 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")
+            
+            # 임시 파일 정리
+            try:
+                os.unlink(original_path)
+                if background_path:
+                    os.unlink(background_path)
+                if style_path:
+                    os.unlink(style_path)
+            except Exception as e:
+                logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
+            
             return [result_img]
         else:
             logger.error("generate 함수에서 None 반환됨.")
@@ -212,9 +210,9 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
-            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")
+            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")
             prompt_input = gr.Textbox(
                 lines=2,
                 placeholder="편집할 내용을 입력하세요...",
@@ -223,6 +221,7 @@ with gr.Blocks() as demo:
             submit_btn = gr.Button("이미지 편집 실행")
         with gr.Column():
             output_gallery = gr.Gallery(label="편집 결과")
+            output_text = gr.Textbox(label="API 응답 텍스트", visible=True)
 
     submit_btn.click(
         fn=process_image_and_prompt,
@@ -230,19 +229,17 @@ with gr.Blocks() as demo:
         outputs=output_gallery,
     )
 
-# --- 테스트 코드 ---
-# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
-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("직접 호출 실패.")
-
-demo.launch(share=True)
+    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>
+        </ul>
+    </div>
+    """)
+
+# --- 실행 ---
+demo.launch(share=True)