Update app.py

app.py

@@ -47,16 +47,31 @@ def merge_images(person_img_path, product_img_path, background_img_path, prompt,
         product_img = Image.open(product_img_path)
         
         # 컨텐츠 리스트 생성
-        contents = [person_img, product_img]
+        contents = []
+        
+        # 사람 이미지 추가
+        contents.append(
+            types.Part.from_data(data=person_img, mime_type="image/jpeg")
+        )
+        
+        # 상품 이미지 추가
+        contents.append(
+            types.Part.from_data(data=product_img, mime_type="image/jpeg")
+        )
         
         # 배경 이미지가 있으면 추가
         if background_img_path:
             background_img = Image.open(background_img_path)
-            contents.append(background_img)
+            contents.append(
+                types.Part.from_data(data=background_img, mime_type="image/jpeg")
+            )
             logger.debug("배경 이미지 추가됨")
         
         # 마지막에 프롬프트 추가
-        contents.append(prompt)
+        contents.append(
+            types.Part.from_text(text=prompt)
+        )
+        
         logger.debug(f"컨텐츠 객체 생성 완료: {len(contents)} 아이템")
 
         # 생성 설정
@@ -65,7 +80,7 @@ def merge_images(person_img_path, product_img_path, background_img_path, prompt,
             top_p=0.95,
             top_k=40,
             max_output_tokens=8192,
-            response_modalities=["text", "image"],
+            response_modalities=["image", "text"],
         )
         logger.debug(f"생성 설정: {generate_content_config}")
 
@@ -73,43 +88,70 @@ def merge_images(person_img_path, product_img_path, background_img_path, prompt,
             temp_path = tmp.name
             logger.debug(f"임시 파일 생성됨: {temp_path}")
 
-            # 단일 요청으로 이미지 생성
-            response = client.models.generate_content(
+            # 스트리밍 방식으로 데이터 수신 (원래 코드 방식으로 복원)
+            response_stream = client.models.generate_content_stream(
                 model=model,
-                contents=contents,
+                contents=[
+                    types.Content(
+                        role="user",
+                        parts=contents,
+                    ),
+                ],
                 config=generate_content_config,
             )
             
-            logger.debug("응답 처리 시작...")
+            logger.debug("응답 스트림 처리 시작...")
             
             # 응답에서 이미지와 텍스트 추출
             image_saved = False
             response_text = ""
             
-            for part in response.candidates[0].content.parts:
-                if hasattr(part, 'text') and part.text:
-                    response_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)
-                    logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
-                    image_saved = True
+            for chunk in response_stream:
+                logger.debug(f"chunk 수신: {chunk}")
+                
+                # 응답 검증
+                if not hasattr(chunk, 'candidates') or not chunk.candidates:
+                    logger.warning("chunk에 candidates가 없습니다. 건너뜁니다.")
+                    continue
+                    
+                if not hasattr(chunk.candidates[0], 'content') or not chunk.candidates[0].content:
+                    logger.warning("chunk.candidates[0]에 content가 없습니다. 건너뜁니다.")
+                    continue
+                    
+                if not hasattr(chunk.candidates[0].content, 'parts') or not chunk.candidates[0].content.parts:
+                    logger.warning("chunk.candidates[0].content에 parts가 없습니다. 건너뜁니다.")
+                    continue
+                
+                for part in chunk.candidates[0].content.parts:
+                    if hasattr(part, 'text') and part.text:
+                        response_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)
+                        logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
+                        image_saved = True
 
             if not image_saved:
                 logger.warning("이미지가 생성되지 않았습니다.")
-                return None, response_text
+                return None, response_text or "이미지가 생성되지 않았습니다. 다른 프롬프트나 이미지로 시도해보세요."
 
         logger.debug("이미지 생성 완료.")
         return temp_path, response_text
 
     except Exception as e:
         logger.exception("이미지 생성 중 오류 발생:")
-        return None, str(e)  # 오류 발생 시 None과 오류 메시지 반환
+        return None, f"오류 발생: {str(e)}"  # 오류 발생 시 None과 오류 메시지 반환
 
 
 def process_images_and_prompt(person_pil, product_pil, background_pil, prompt):
     logger.debug(f"process_images_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
     try:
+        # 필수 이미지 확인
+        if person_pil is None:
+            return [], "사람 이미지를 업로드해주세요."
+        if product_pil is None:
+            return [], "상품 이미지를 업로드해주세요."
+            
         # 기본 프롬프트 설정 (비어있는 경우)
         if not prompt or not prompt.strip():
             if background_pil:
@@ -125,22 +167,34 @@ def process_images_and_prompt(person_pil, product_pil, background_pil, prompt):
                 prompt += " Create a realistic composite image of this person with this product."
         
         # 이미지 저장
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_person:
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_person:
             person_path = tmp_person.name
-            person_pil.save(person_path)
+            # RGB로 변환하여 저장
+            person_img = person_pil
+            if person_img.mode != 'RGB':
+                person_img = person_img.convert('RGB')
+            person_img.save(person_path, 'JPEG')
             logger.debug(f"사람 이미지 저장 완료: {person_path}")
             
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_product:
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_product:
             product_path = tmp_product.name
-            product_pil.save(product_path)
+            # RGB로 변환하여 저장
+            product_img = product_pil
+            if product_img.mode != 'RGB':
+                product_img = product_img.convert('RGB')
+            product_img.save(product_path, 'JPEG')
             logger.debug(f"상품 이미지 저장 완료: {product_path}")
         
         # 배경 이미지 저장 (있는 경우)
         background_path = None
         if background_pil is not None:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_bg:
+            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_bg:
                 background_path = tmp_bg.name
-                background_pil.save(background_path)
+                # RGB로 변환하여 저장
+                background_img = background_pil
+                if background_img.mode != 'RGB':
+                    background_img = background_img.convert('RGB')
+                background_img.save(background_path, 'JPEG')
                 logger.debug(f"배경 이미지 저장 완료: {background_path}")
 
         # 이미지 합성 실행
@@ -154,27 +208,31 @@ def process_images_and_prompt(person_pil, product_pil, background_pil, prompt):
         # 이미지 반환 및 임시 파일 정리
         if result_path:
             logger.debug(f"이미지 생성 완료. 경로: {result_path}")
-            result_img = Image.open(result_path)
-            if result_img.mode == "RGBA":
-                result_img = result_img.convert("RGB")
-            
-            # 임시 파일 정리
             try:
-                os.unlink(person_path)
-                os.unlink(product_path)
-                if background_path:
-                    os.unlink(background_path)
+                result_img = Image.open(result_path)
+                if result_img.mode == "RGBA":
+                    result_img = result_img.convert("RGB")
+                
+                # 임시 파일 정리
+                try:
+                    os.unlink(person_path)
+                    os.unlink(product_path)
+                    if background_path:
+                        os.unlink(background_path)
+                except Exception as e:
+                    logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
+                
+                return [result_img], response_text or "이미지가 성공적으로 생성되었습니다."
             except Exception as e:
-                logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
-            
-            return [result_img], response_text
+                logger.exception(f"결과 이미지 로드 중 오류: {str(e)}")
+                return [], f"결과 이미지 처리 중 오류: {str(e)}"
         else:
             logger.error("merge_images 함수에서 None 반환됨.")
-            return [], response_text  # 오류 시 빈 리스트 반환
+            return [], response_text or "이미지 생성에 실패했습니다. 다른 이미지나 프롬프트로 시도해보세요."
 
     except Exception as e:
         logger.exception("process_images_and_prompt 함수에서 오류 발생:")
-        return [], str(e)  # 오류 �� 빈 리스트와 오류 메시지 반환
+        return [], f"오류 발생: {str(e)}"  # 오류 시 빈 리스트와 오류 메시지 반환
 
 
 # --- Gradio 인터페이스 구성 ---