Update app.py

app.py

@@ -12,7 +12,7 @@ import logging
 from google import genai
 from google.genai import types
 
-# .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
+# .env 파일에 저장된 환경변수 로드 (pip install python-dotenv 필요)
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -29,11 +29,11 @@ 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"):
-    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일명: '{file_name}', 모델: '{model}'")
-
+def generate(text, composite_file, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
+    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', composite_file: '{composite_file}', "
+                 f"background_file: '{background_file}', style_file: '{style_file}', 모델: '{model}'")
     try:
-        # API 키는 환경변수에서 불러옴
+        # 환경변수에서 API 키 불러오기
         effective_api_key = os.environ.get("GEMINI_API_KEY")
         if effective_api_key:
             logger.debug("환경변수에서 API 키 불러옴")
@@ -44,25 +44,32 @@ 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}")
-
-        # 컨텐츠 객체 생성: 파일 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),
-                ],
-            ),
-        ]
+        # 업로드할 파일들 리스트 생성
+        files = []
+        # 원본 이미지 업로드 (필수)
+        composite_upload = client.files.upload(file=composite_file)
+        files.append(composite_upload)
+        logger.debug(f"원본 이미지 업로드 완료. URI: {composite_upload.uri}, MIME 타입: {composite_upload.mime_type}")
+
+        # 배경 이미지 업로드 (선택)
+        if background_file is not None:
+            background_upload = client.files.upload(file=background_file)
+            files.append(background_upload)
+            logger.debug(f"배경 이미지 업로드 완료. URI: {background_upload.uri}, MIME 타입: {background_upload.mime_type}")
+        # 스타일 이미지 업로드 (선택)
+        if style_file is not None:
+            style_upload = client.files.upload(file=style_file)
+            files.append(style_upload)
+            logger.debug(f"스타일 이미지 업로드 완료. URI: {style_upload.uri}, MIME 타입: {style_upload.mime_type}")
+
+        # 파일 업로드 결과를 바탕으로 콘텐츠 파트 생성
+        parts = []
+        for file_obj in files:
+            parts.append(types.Part.from_uri(file_uri=file_obj.uri, mime_type=file_obj.mime_type))
+        # 마지막에 텍스트 프롬프트 추가
+        parts.append(types.Part.from_text(text=text))
+
+        contents = [types.Content(role="user", parts=parts)]
         logger.debug(f"컨텐츠 객체 생성 완료: {contents}")
 
         generate_content_config = types.GenerateContentConfig(
@@ -70,10 +77,7 @@ def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
             top_p=0.95,
             top_k=40,
             max_output_tokens=8192,
-            response_modalities=[
-                "image",
-                "text",
-            ],
+            response_modalities=["image", "text"],
             response_mime_type="text/plain",
         )
         logger.debug(f"생성 설정: {generate_content_config}")
@@ -113,19 +117,38 @@ 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(composite_pil, background_pil, style_pil, prompt):
     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}")
+            logger.debug(f"원본 이미지 저장 완료: {composite_path}")
 
-        file_name = composite_path
-        input_text = prompt
-        model = "gemini-2.0-flash-exp-image-generation"
+        background_path = None
+        style_path = None
+
+        # 배경 이미지 저장 (입력된 경우)
+        if background_pil is not None:
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp2:
+                background_path = tmp2.name
+                background_pil.save(background_path)
+                logger.debug(f"배경 이미지 저장 완료: {background_path}")
 
-        gemma_edited_image_path = generate(text=input_text, file_name=file_name, model=model)
+        # 스타일 이미지 저장 (입력된 경우)
+        if style_pil is not None:
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp3:
+                style_path = tmp3.name
+                style_pil.save(style_path)
+                logger.debug(f"스타일 이미지 저장 완료: {style_path}")
+
+        model = "gemini-2.0-flash-exp-image-generation"
+        gemma_edited_image_path = generate(text=prompt,
+                                           composite_file=composite_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 +180,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")
+            composite_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 +198,18 @@ with gr.Blocks() as demo:
 
     submit_btn.click(
         fn=process_image_and_prompt,
-        inputs=[image_input, prompt_input],
+        inputs=[composite_input, background_input, style_input, prompt_input],
         outputs=output_gallery,
     )
 
 # --- 테스트 코드 ---
-# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
-dummy_image = Image.new("RGBA", (100, 100), color="red")
-dummy_prompt = "이미지를 파란색으로 변경해줘"
+dummy_composite = 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_composite, dummy_background, dummy_style, dummy_prompt)
 
 if result:
     logger.info(f"직접 호출 성공. 결과: {result}")