Update app.py

app.py

@@ -4,6 +4,8 @@ from PIL import Image
 import gradio as gr
 import logging
 import re
+import io
+from io import BytesIO
 
 from google import genai
 from google.genai import types
@@ -92,13 +94,14 @@ def preprocess_prompt(prompt, image1, image2, image3):
     elif "을 붉은색으로 바꿔라" in prompt or "를 붉은색으로 바꿔라" in prompt:
         prompt = "첫 번째 이미지를 붉은색 톤으로 변경해주세요. 전체적인 색상을 붉은 계열로 조정하고 자연스러운 느낌을 유지해주세요."
     
-    # 입력이 위 기능들에 해당하지 않는 경우 원본 프롬프트 유지
+    # 명확한 이미지 생성 요청 추가
+    prompt += " 이미지를 생성해주세요."
     
     return prompt
 
 def generate_with_images(prompt, images):
     """
-    새로운 API 호출 방식으로 구현 - 응답 모달리티 수정
+    공식 문서에 기반한 올바른 API 호출 방식 구현
     """
     try:
         # API 키 확인
@@ -111,138 +114,57 @@ def generate_with_images(prompt, images):
         
         logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
         
-        # 임시 파일 생성 및 처리
-        temp_image_paths = []
-        file_uris = []
+        # 컨텐츠 준비
+        contents = []
         
-        # 이미지 파일 저장 및 업로드
+        # 텍스트 프롬프트 추가
+        contents.append(prompt)
+        
+        # 이미지 추가
         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}")
+                contents.append(img)
+                logger.info(f"이미지 #{idx} 추가됨")
         
-        # 컨텐츠 파트 구성
-        parts = []
-        
-        # 이미지 파트 추가
-        for file in file_uris:
-            parts.append(
-                types.Part.from_uri(
-                    file_uri=file.uri,
-                    mime_type=file.mime_type,
-                )
+        # 생성 설정 - 공식 문서에 따라 responseModalities 설정
+        response = client.models.generate_content(
+            model="gemini-2.0-flash-exp-image-generation",
+            contents=contents,
+            config=types.GenerateContentConfig(
+                response_modalities=['Text', 'Image'],
+                temperature=1,
+                top_p=0.95,
+                top_k=40,
+                max_output_tokens=8192
             )
-        
-        # 텍스트 프롬프트 추가
-        parts.append(types.Part.from_text(text=prompt))
-        
-        # 요청 메시지 구성
-        contents = [
-            types.Content(
-                role="user",
-                parts=parts,
-            ),
-        ]
-        
-        # 생성 설정 - response_modalities 제거
-        generate_content_config = types.GenerateContentConfig(
-            temperature=1,
-            top_p=0.95,
-            top_k=40,
-            max_output_tokens=8192,
         )
         
         # 임시 파일 생성
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
             temp_path = tmp.name
             
-            # Gemini 모델로 요청 전송
-            response = client.models.generate_content(
-                model="gemini-pro-vision",  # 모델을 gemini-pro-vision으로 변경
-                contents=contents,
-                config=generate_content_config,
-            )
+            result_text = ""
+            image_found = False
             
-            # 응답이 이미지 대신 텍스트일 것으로 예상
-            response_text = ""
+            # 응답 처리
             for part in response.candidates[0].content.parts:
                 if hasattr(part, 'text') and part.text:
-                    response_text += part.text
+                    result_text += part.text
+                    logger.info(f"응답 텍스트: {part.text}")
+                elif hasattr(part, 'inline_data') and part.inline_data:
+                    save_binary_file(temp_path, part.inline_data.data)
+                    image_found = True
+                    logger.info("응답에서 이미지 추출 성공")
             
-            logger.info(f"응답 텍스트: {response_text}")
-            
-            # 이제 이 텍스트 응답을 사용하여 gemini-2.0-flash-exp-image-generation 모델로 이미지 생성
-            # 첫 번째 이미지만 사용 (여러 이미지 사용은 복잡해질 수 있음)
-            if len(temp_image_paths) > 0:
-                with open(temp_image_paths[0], 'rb') as f:
-                    image_data = f.read()
-                
-                # gemini-2.0-flash-exp-image-generation 모델을 사용한 단일 이미지 요청
-                response_stream = client.models.generate_content_stream(
-                    model="gemini-2.0-flash-exp-image-generation",
-                    contents=[
-                        types.Content(
-                            role="user",
-                            parts=[
-                                types.Part.from_data(
-                                    mime_type="image/png",
-                                    data=image_data
-                                ),
-                                types.Part.from_text(text=f"{prompt} {response_text}")
-                            ]
-                        )
-                    ],
-                    config=types.GenerateContentConfig(
-                        temperature=1,
-                        response_mime_type="image/png",
-                    )
-                )
-                
-                # 응답에서 이미지 추출
-                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
+            if not image_found:
+                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
                     
-                    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("응답에서 이미지 추출 성공")
+            # 결과 이미지 반환
+            result_img = Image.open(temp_path)
+            if result_img.mode == "RGBA":
+                result_img = result_img.convert("RGB")
                 
-                # 임시 파일 정리
-                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에서 이미지를 생성하지 못했습니다. 다른 프롬프트로 시도해보세요."
-                        
-                # 결과 이미지 반환
-                result_img = Image.open(temp_path)
-                if result_img.mode == "RGBA":
-                    result_img = result_img.convert("RGB")
-                    
-                return result_img, "이미지가 성공적으로 생성되었습니다."
-            else:
-                return None, "처리할 이미지가 없습니다."
+            return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
             
     except Exception as e:
         logger.exception("이미지 생성 중 오류 발생:")