Update app.py

app.py

@@ -4,7 +4,6 @@ from PIL import Image
 import gradio as gr
 import logging
 import re
-import io
 
 from google import genai
 from google.genai import types
@@ -25,6 +24,7 @@ def preprocess_prompt(prompt, image1, image2, image3):
     """
     프롬프트를 처리하고 기능 명령을 해석
     """
+    # 기존 preprocess_prompt 함수 코드 유지
     # 이미지 없는 참조 확인 및 처리
     has_img1 = image1 is not None
     has_img2 = image2 is not None
@@ -96,9 +96,9 @@ def preprocess_prompt(prompt, image1, image2, image3):
     
     return prompt
 
-def process_images_with_prompt(image1, image2, image3, prompt):
+def generate_with_images(prompt, images):
     """
-    3개의 이미지와 프롬프트를 처리하는 함수
+    문제가 된 API 호출 부분을 새로운 방식으로 구현
     """
     try:
         # API 키 확인
@@ -109,86 +109,91 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         # Gemini 클라이언트 초기화
         client = genai.Client(api_key=api_key)
         
-        # 이미지 개수 확인
-        images = [image1, image2, image3]
-        valid_images = [img for img in images if img is not None]
+        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
         
-        if not valid_images:
-            return None, "적어도 하나의 이미지를 업로드해주세요."
+        # 임시 파일 생성 및 처리
+        temp_image_paths = []
+        file_uris = []
         
-        # 프롬프트 처리
-        if not prompt or not prompt.strip():
-            # 프롬프트가 없으면 업로드된 이미지 수에 따라 자동 합성 프롬프트 생성
-            if len(valid_images) == 1:
-                prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
-                logger.info("단일 이미지 프롬프트 자동 생성")
-            elif len(valid_images) == 2:
-                prompt = "이 두 이미지를 자연스럽게 합성해주세요. 두 이미지의 요소를 조화롭게 통합하여 하나의 이미지로 만들어주세요."
-                logger.info("두 이미지 합성 프롬프트 자동 생성")
-            else:
-                prompt = "이 세 이미지를 창의적으로 합성해주세요. 모든 이미지의 주요 요소를 포함하되 자연스럽고 일관된 하나의 장면으로 만들어주세요."
-                logger.info("세 이미지 합성 프롬프트 자동 생성")
-        else:
-            # 프롬프트 전처리 및 기능 명령 해석
-            prompt = preprocess_prompt(prompt, image1, image2, image3)
+        # 이미지 파일 저장 및 업로드
+        for idx, img in enumerate(images, 1):
+            if img is not None:
+                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+                    temp_path = tmp.name
+                    img.save(temp_path, format="PNG")
+                    temp_image_paths.append(temp_path)
+                    
+                    # 파일 업로드
+                    uploaded_file = client.files.upload(file=temp_path)
+                    file_uris.append(uploaded_file)
+                    logger.info(f"이미지 #{idx} 업로드됨: {uploaded_file.uri}")
         
-        # 컨텐츠 리스트 생성 (이미지와 프롬프트 결합)
+        # 컨텐츠 파트 구성
         parts = []
         
+        # 이미지 파트 추가
+        for file in file_uris:
+            parts.append(
+                types.Part.from_uri(
+                    file_uri=file.uri,
+                    mime_type=file.mime_type,
+                )
+            )
+        
         # 텍스트 프롬프트 추가
-        parts.append({
-            "text": prompt
-        })
+        parts.append(types.Part.from_text(text=prompt))
         
-        # 사용 가능한 이미지 추가
-        for idx, img in enumerate(images, 1):
-            if img is not None:
-                # PIL 이미지를 바이트로 변환
-                img_byte_arr = io.BytesIO()
-                img.save(img_byte_arr, format='PNG')
-                img_bytes = img_byte_arr.getvalue()
-                
-                # 이미지를 파트로 추가 (genai.types.Part 대신 직접 딕셔너리 사용)
-                parts.append({
-                    "inline_data": {
-                        "mime_type": "image/png",
-                        "data": img_bytes
-                    }
-                })
-                logger.info(f"이미지 #{idx} 추가됨")
+        # 요청 메시지 구성
+        contents = [
+            types.Content(
+                role="user",
+                parts=parts,
+            ),
+        ]
         
         # 생성 설정
-        generate_content_config = {
-            "temperature": 1.0,
-            "response_modalities": ["image"]
-        }
+        generate_content_config = types.GenerateContentConfig(
+            temperature=1,
+            response_modalities=["image"],
+        )
         
         # 임시 파일 생성
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
             temp_path = tmp.name
             
             # Gemini 모델로 요청 전송
-            logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
-            
-            # API 요청 형식 수정
-            response = client.models.generate_content(
+            response_stream = client.models.generate_content_stream(
                 model="gemini-2.0-flash-exp-image-generation",
-                contents=[{
-                    "role": "user",
-                    "parts": parts
-                }],
-                generation_config=generate_content_config
+                contents=contents,
+                config=generate_content_config,
             )
             
             # 응답에서 이미지 추출
             image_found = False
+            for chunk in response_stream:
+                if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
+                    logger.warning("청크에 데이터가 없음. 건너뜀.")
+                    continue
+                
+                for part in chunk.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("응답에서 이미지 추출 성공")
             
-            # 응답 구조 처리
-            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("응답에서 이미지 추출 성공")
+            # 임시 파일 정리
+            for path in temp_image_paths:
+                try:
+                    os.unlink(path)
+                except:
+                    pass
+            
+            # 파일 URI 정리
+            for file in file_uris:
+                try:
+                    client.files.delete(fileLocation=file.uri)
+                except:
+                    pass
             
             if not image_found:
                 return None, "API에서 이미지를 생성하지 못했습니다. 다른 프롬프트로 시도해보세요."
@@ -204,7 +209,42 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         logger.exception("이미지 생성 중 오류 발생:")
         return None, f"오류 발생: {str(e)}"
 
-# 기능 선택 콜백
+def process_images_with_prompt(image1, image2, image3, prompt):
+    """
+    3개의 이미지와 프롬프트를 처리하는 함수
+    """
+    try:
+        # 이미지 개수 확인
+        images = [image1, image2, image3]
+        valid_images = [img for img in images if img is not None]
+        
+        if not valid_images:
+            return None, "적어도 하나의 이미지를 업로드해주세요."
+        
+        # 프롬프트 처리
+        if not prompt or not prompt.strip():
+            # 프롬프트가 없으면 업로드된 이미지 수에 따라 자동 합성 프롬프트 생성
+            if len(valid_images) == 1:
+                prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
+                logger.info("단일 이미지 프롬프트 자동 생성")
+            elif len(valid_images) == 2:
+                prompt = "이 두 이미지를 자연스럽게 합성해주세요. 두 이미지의 요소를 조화롭게 통합하여 하나의 이미지로 만들어주세요."
+                logger.info("두 이미지 합성 프롬프트 자동 생성")
+            else:
+                prompt = "이 세 이미지를 창의적으로 합성해주세요. 모든 이미지의 주요 요소를 포함하되 자연스럽고 일관된 하나의 장면으로 만들어주세요."
+                logger.info("세 이미지 합성 프롬프트 자동 생성")
+        else:
+            # 프롬프트 전처리 및 기능 명령 해석
+            prompt = preprocess_prompt(prompt, image1, image2, image3)
+        
+        # 새로운 API 호출 방식 사용
+        return generate_with_images(prompt, valid_images)
+            
+    except Exception as e:
+        logger.exception("이미지 처리 중 오류 발생:")
+        return None, f"오류 발생: {str(e)}"
+
+# 기능 선택 콜백 (기존 코드 유지)
 def update_prompt_from_function(function_choice, custom_text=""):
     function_templates = {
         "1. 이미지 변경": f'#1을 "{custom_text if custom_text else "원하는 설명"}"으로 바꿔라',
@@ -218,7 +258,7 @@ def update_prompt_from_function(function_choice, custom_text=""):
     
     return function_templates.get(function_choice, "")
 
-# Gradio 인터페이스
+# Gradio 인터페이스 (기존 ���드 유지)
 with gr.Blocks() as demo:
     gr.HTML(
         """