Update app.py

app.py

@@ -9,9 +9,8 @@ import base64
 import mimetypes
 import logging
 
-# 기존의 "from google import genai" 대신 다음과 같이 수정
-import google.generativeai as genai
-from google.generativeai import types
+from google import genai
+from google.genai import types
 
 # .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
 from dotenv import load_dotenv
@@ -22,14 +21,17 @@ logging.basicConfig(level=logging.DEBUG,
                     format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+
 def save_binary_file(file_name, data):
     logger.debug(f"파일에 이진 데이터 저장 중: {file_name}")
     with open(file_name, "wb") as f:
         f.write(data)
     logger.debug(f"파일 저장 완료: {file_name}")
 
-def generate(text, file_name, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
+
+def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
     logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일명: '{file_name}', 모델: '{model}'")
+
     try:
         # API 키는 환경변수에서 불러옴
         effective_api_key = os.environ.get("GEMINI_API_KEY")
@@ -42,53 +44,23 @@ def generate(text, file_name, background_file=None, style_file=None, model="gemi
         client = genai.Client(api_key=effective_api_key)
         logger.debug("Gemini 클라이언트 초기화 완료.")
 
-        # 원본, 배경, 스타일 이미지를 각각 업로드 (존재하는 경우)
-        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:
-            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))
+        # 파일 업로드
+        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=parts,
+                parts=[
+                    types.Part.from_uri(
+                        file_uri=files[0].uri,
+                        mime_type=files[0].mime_type,
+                    ),
+                    types.Part.from_text(text=text),
+                ],
             ),
         ]
         logger.debug(f"컨텐츠 객체 생성 완료: {contents}")
@@ -132,47 +104,28 @@ def generate(text, file_name, background_file=None, style_file=None, model="gemi
 
                 logger.debug(f"Raw chunk: {chunk}")
 
-        del uploaded_files
+        del files
         logger.debug("업로드된 파일 정보 삭제 완료.")
         return temp_path
 
     except Exception as e:
         logger.exception("이미지 생성 중 오류 발생:")
-        return None
+        return None  # 오류 발생 시 None 반환
+
 
-def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pil=None):
+def process_image_and_prompt(composite_pil, prompt):
     logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
     try:
-        # 원본 이미지 저장
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            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}")
+            composite_path = tmp.name
+            composite_pil.save(composite_path)
+            logger.debug(f"합성 이미지 저장 완료: {composite_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=original_path,
-                                           background_file=background_path,
-                                           style_file=style_path,
-                                           model=model)
+        gemma_edited_image_path = generate(text=input_text, file_name=file_name, model=model)
 
         if gemma_edited_image_path:
             logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
@@ -182,10 +135,12 @@ def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pi
             return [result_img]
         else:
             logger.error("generate 함수에서 None 반환됨.")
-            return []
+            return []  # 오류 시 빈 리스트 반환
+
     except Exception as e:
         logger.exception("process_image_and_prompt 함수에서 오류 발생:")
-        return []
+        return []  # 오류 시 빈 리스트 반환
+
 
 # --- Gradio 인터페이스 구성 ---
 with gr.Blocks() as demo:
@@ -202,13 +157,11 @@ with gr.Blocks() as demo:
         </div>
         """
     )
-    gr.Markdown("원본, 배경, 스타일 이미지를 업로드하고, 편집할 내용을 입력하세요.")
+    gr.Markdown("이미지를 업로드하고, 편집할 내용을 입력하세요.")
 
     with gr.Row():
         with gr.Column():
-            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")
+            image_input = gr.Image(type="pil", label="이미지 업로드", image_mode="RGBA")
             prompt_input = gr.Textbox(
                 lines=2,
                 placeholder="편집할 내용을 입력하세요...",
@@ -220,18 +173,17 @@ with gr.Blocks() as demo:
 
     submit_btn.click(
         fn=process_image_and_prompt,
-        inputs=[original_input, prompt_input, background_input, style_input],
+        inputs=[image_input, prompt_input],
         outputs=output_gallery,
     )
 
 # --- 테스트 코드 ---
-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 = "이미지 편집: 원본은 빨강, 배경은 초록, 스타일은 파랑"
+# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
+dummy_image = Image.new("RGBA", (100, 100), color="red")
+dummy_prompt = "이미지를 파란색으로 변경해줘"
 
 logger.info("process_image_and_prompt 함수를 직접 호출합니다...")
-result = process_image_and_prompt(dummy_original, dummy_prompt, dummy_background, dummy_style)
+result = process_image_and_prompt(dummy_image, dummy_prompt)
 
 if result:
     logger.info(f"직접 호출 성공. 결과: {result}")