Update app.py

app.py

@@ -156,43 +156,74 @@ def preprocess_prompt(prompt, image1, image2, image3):
     prompt += " 이미지를 생성해주세요. 이미지에 텍스트나 글자를 포함하지 마세요."
     return prompt
 
-def convert_to_jpg(pil_image):
-    """이미지를 JPG 형식으로 변환합니다."""
-    import tempfile
-    from PIL import Image
-    import io
-    
-    if pil_image is None:
-        return None
-        
-    # RGBA 이미지를 RGB로 변환 (JPG는 알파 채널을 지원하지 않음)
-    if pil_image.mode == "RGBA":
-        pil_image = pil_image.convert("RGB")
-    
-    # 이미지를 JPG로 변환
-    jpg_buffer = io.BytesIO()
-    pil_image.save(jpg_buffer, format="JPEG", quality=95)
-    jpg_buffer.seek(0)
-    
-    # JPG 이미지 다시 로드
-    jpg_image = Image.open(jpg_buffer)
-    return jpg_image
+def generate_with_images(prompt, images, variation_index=0):
+    # 기존 함수 유지
+    try:
+        api_key = os.environ.get("GEMINI_API_KEY")
+        if not api_key:
+            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
 
-def save_as_jpg(img):
-    """이미지를 강제로 JPG로 저장하고 그 경로를 반환합니다."""
-    if img is None:
-        return None
-        
-    # RGB로 변환 (JPG는 알파 채널을 지원하지 않음)
-    if img.mode == "RGBA":
-        img = img.convert("RGB")
+        client = genai.Client(api_key=api_key)
+        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")
+
+        variation_suffixes = [
+            " Create this as the first variation. Do not add any text, watermarks, or labels to the image.",
+            " Create this as the second variation with more vivid colors. Do not add any text, watermarks, or labels to the image.",
+            " Create this as the third variation with a more creative style. Do not add any text, watermarks, or labels to the image.",
+            " Create this as the fourth variation with enhanced details. Do not add any text, watermarks, or labels to the image."
+        ]
         
-    # 임시 파일로 JPG 저장
-    temp_jpg = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
-    img.save(temp_jpg.name, format="JPEG", quality=95)
-    temp_jpg.close()
-    
-    return temp_jpg.name
+        if variation_index < len(variation_suffixes):
+            prompt = prompt + variation_suffixes[variation_index]
+        else:
+            prompt = prompt + " Do not add any text, watermarks, or labels to the image."
+
+        contents = [prompt]
+        for idx, img in enumerate(images, 1):
+            if img is not None:
+                contents.append(img)
+                logger.info(f"이미지 #{idx} 추가됨")
+
+        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
+            )
+        )
+
+        # 임시 파일은 항상 JPG 확장자로 생성
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+            temp_path = tmp.name
+            result_text = ""
+            image_found = False
+            for part in response.candidates[0].content.parts:
+                if hasattr(part, 'text') and 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("응답에서 이미지 추출 성공")
+            if not image_found:
+                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
+            
+            result_img = Image.open(temp_path)
+            if result_img.mode == "RGBA":
+                result_img = result_img.convert("RGB")  # JPG는 투명도를 지원하지 않으므로 RGB로 변환
+            
+            # 변환된 이미지를 JPG로 저장
+            result_img.save(temp_path, format="JPEG", quality=95)
+            
+            # 파일 경로 반환
+            return temp_path, f"이미지가 성공적으로 생성되었습니다. {result_text}"
+    except Exception as e:
+        logger.exception("이미지 생성 중 오류 발생:")
+        return None, f"오류 발생: {str(e)}"
 
 def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0, max_retries=3):
     # 기존 함수 내용 유지
@@ -225,8 +256,7 @@ def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0
 
             result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
             if result_img is not None:
-                # 이미지를 JPG로 변환하여 반환
-                result_img = convert_to_jpg(result_img)
+                # 이미 파일 경로가 반환되므로 추가 처리 없이 그대로 반환
                 return result_img, status, final_prompt
             else:
                 last_error = status
@@ -257,8 +287,7 @@ def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progres
         result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, i, max_retries)
         
         if result_img is not None:
-            # 이미지를 JPG로 변환
-            result_img = convert_to_jpg(result_img)
+            # 이미 파일 경로이므로 추가 처리 없이 그대로 추가
             results.append(result_img)
             statuses.append(f"이미지 #{i+1}: {status}")
             prompts.append(f"이미지 #{i+1}: {final_prompt}")