Update app.py

app.py

@@ -29,7 +29,7 @@ def save_binary_file(file_name, data):
     logger.debug(f"파일 저장 완료: {file_name}")
 
 
-def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
+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}'")
 
     try:
@@ -44,23 +44,54 @@ def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
         client = genai.Client(api_key=effective_api_key)
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
-        # 파일 업로드
-        files = [
-            client.files.upload(file=file_name),
-        ]
-        logger.debug(f"파일 업로드 완료. URI: {files[0].uri}, MIME 타입: {files[0].mime_type}")
+        # 원본, 배경, 스타일 이미지를 각각 업로드 (존재하는 경우)
+        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))
 
-        # 컨텐츠 객체 생성: 파일 URI와 텍스트 프롬프트를 함께 포함
         contents = [
             types.Content(
                 role="user",
-                parts=[
-                    types.Part.from_uri(
-                        file_uri=files[0].uri,
-                        mime_type=files[0].mime_type,
-                    ),
-                    types.Part.from_text(text=text),
-                ],
+                parts=parts,
             ),
         ]
         logger.debug(f"컨텐츠 객체 생성 완료: {contents}")
@@ -104,7 +135,7 @@ def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
 
                 logger.debug(f"Raw chunk: {chunk}")
 
-        del files
+        del uploaded_files
         logger.debug("업로드된 파일 정보 삭제 완료.")
         return temp_path
 
@@ -113,19 +144,39 @@ def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
         return None  # 오류 발생 시 None 반환
 
 
-def process_image_and_prompt(composite_pil, prompt):
+def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pil=None):
     logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
     try:
+        # 원본 이미지 저장
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            composite_path = tmp.name
-            composite_pil.save(composite_path)
-            logger.debug(f"합성 이미지 저장 완료: {composite_path}")
+            original_path = tmp.name
+            original_pil.save(original_path)
+            logger.debug(f"원본 이미지 저장 완료: {original_path}")
+
+        # 배경 이미지 저장 (제공된 경우)
+        background_path = None
+        if background_pil is not None:
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_bg:
+                background_path = tmp_bg.name
+                background_pil.save(background_path)
+                logger.debug(f"배경 이미지 저장 ��료: {background_path}")
+
+        # 스타일 이미지 저장 (제공된 경우)
+        style_path = None
+        if style_pil is not None:
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_style:
+                style_path = tmp_style.name
+                style_pil.save(style_path)
+                logger.debug(f"스타일 이미지 저장 완료: {style_path}")
 
-        file_name = composite_path
         input_text = prompt
         model = "gemini-2.0-flash-exp-image-generation"
 
-        gemma_edited_image_path = generate(text=input_text, file_name=file_name, model=model)
+        gemma_edited_image_path = generate(text=input_text,
+                                           file_name=original_path,
+                                           background_file=background_path,
+                                           style_file=style_path,
+                                           model=model)
 
         if gemma_edited_image_path:
             logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
@@ -157,11 +208,13 @@ with gr.Blocks() as demo:
         </div>
         """
     )
-    gr.Markdown("이미지를 업로드하고, 편집할 내용을 입력하세요.")
+    gr.Markdown("원본, 배경, 스타일 이미지를 업로드하고, 편집할 내용을 입력하세요.")
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="이미지 업로드", image_mode="RGBA")
+            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")
             prompt_input = gr.Textbox(
                 lines=2,
                 placeholder="편집할 내용을 입력하세요...",
@@ -173,17 +226,19 @@ with gr.Blocks() as demo:
 
     submit_btn.click(
         fn=process_image_and_prompt,
-        inputs=[image_input, prompt_input],
+        inputs=[original_input, prompt_input, background_input, style_input],
         outputs=output_gallery,
     )
 
 # --- 테스트 코드 ---
 # 테스트용 더미 이미지 (실제 이미지로 대체 가능)
-dummy_image = Image.new("RGBA", (100, 100), color="red")
-dummy_prompt = "이미지를 파란색으로 변경해줘"
+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_image, dummy_prompt)
+result = process_image_and_prompt(dummy_original, dummy_prompt, dummy_background, dummy_style)
 
 if result:
     logger.info(f"직접 호출 성공. 결과: {result}")