Update app.py

app.py

@@ -26,41 +26,40 @@ def preprocess_prompt(prompt, image1, image2, image3):
     """
     프롬프트를 처리하고 기능 명령을 해석
     """
-    # 이미지 없는 참조 확인 및 처리
+    # 이미지 존재 여부 확인
     has_img1 = image1 is not None
     has_img2 = image2 is not None
     has_img3 = image3 is not None
     
     # #1, #2, #3 참조를 설명으로 변환 (이미지가 없는 경우 무시)
-    if "#1" in prompt and not has_img1:
-        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
-    else:
-        prompt = prompt.replace("#1", "첫 번째 이미지")
-        
-    if "#2" in prompt and not has_img2:
-        prompt = prompt.replace("#2", "두 번째 이미지(없음)")
-    else:
-        prompt = prompt.replace("#2", "두 번째 이미지")
-        
-    if "#3" in prompt and not has_img3:
-        prompt = prompt.replace("#3", "세 번째 이미지(없음)")
-    else:
-        prompt = prompt.replace("#3", "세 번째 이미지")
+    if "#1" in prompt:
+        prompt = prompt.replace("#1", "첫 번째 이미지" if has_img1 else "첫 번째 이미지(없음)")
+    if "#2" in prompt:
+        prompt = prompt.replace("#2", "두 번째 이미지" if has_img2 else "두 번째 이미지(없음)")
+    if "#3" in prompt:
+        prompt = prompt.replace("#3", "세 번째 이미지" if has_img3 else "세 번째 이미지(없음)")
     
     # 기능 명령 해석
     if "1. 이미지 변경" in prompt:
-        # 기본 프롬프트 제공
-        prompt = "첫 번째 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
+        desc_match = re.search(r'#1을 "(.*?)"으로 바꿔라', prompt)
+        if desc_match:
+            description = desc_match.group(1)
+            prompt = f"첫 번째 이미지를 {description}으로 변경해주세요. 원본 이미지의 주요 내용은 유지하되 새로운 스타일과 분위기로 재해석해주세요."
+        else:
+            prompt = "첫 번째 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
     
     elif "2. 글자지우기" in prompt:
-        # 기본 프롬프트 제공
-        prompt = "첫 번째 이미지에서 모든 텍스트를 찾아 자연스럽게 제거해주세요. 깔끔한 이미지로 만들어주세요."
+        text_match = re.search(r'#1에서 "(.*?)"를 지워라', prompt)
+        if text_match:
+            text_to_remove = text_match.group(1)
+            prompt = f"첫 번째 이미지에서 '{text_to_remove}' 텍스트를 찾아 자연스럽게 제거해주세요. 텍스트가 있던 부분을 배경과 조화롭게 채워주세요."
+        else:
+            prompt = "첫 번째 이미지에서 모든 텍스트를 찾아 자연스럽게 제거해주세요. 깔끔한 이미지로 만들어주세요."
     
     elif "3. 얼굴바꾸기" in prompt:
         prompt = "첫 번째 이미지의 인물 얼굴을 두 번째 이미지의 얼굴로 자연스럽게 교체해주세요. 얼굴의 표정과 특징은 두 번째 이미지를 따르되, 나머지 부분은 첫 번째 이미지를 유지해주세요."
     
     elif "4. 옷바꾸기" in prompt:
-        # 여러 이미지 참조 처리
         if "#3" in prompt or "또는 #3" in prompt:
             prompt = "첫 번째 이미지의 인물 의상을 두 번째 또는 세 번째 이미지의 의상으로 자연스럽게 교체해주세요. 의상의 스타일과 색상은 참조 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."
         else:
@@ -70,7 +69,6 @@ def preprocess_prompt(prompt, image1, image2, image3):
         prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 자연스럽게 교체해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."
     
     elif "6. 이미지 합성(상품포함)" in prompt:
-        # 여러 이미지 참조 처리
         if "#3" in prompt or "또는 #3" in prompt:
             prompt = "첫 번째 이미지와 두 번째, 세 번째 이미지를 자연스럽게 합성해주세요. 모든 이미지의 주요 요소를 포함하고, 특히 상품이 잘 보이도록 조화롭게 통합해주세요."
         else:
@@ -79,151 +77,86 @@ def preprocess_prompt(prompt, image1, image2, image3):
     elif "7. 이미지 합성(스타일적용)" in prompt:
         prompt = "첫 번째 이미지의 내용을 두 번째 이미지의 스타일로 변환해주세요. 첫 번째 이미지의 주요 피사체와 구도는 유지하되, 두 번째 이미지의 예술적 스타일, 색상, 질감을 적용해주세요."
     
-    # 간단한 색상 변경 요청 처리
     elif "을 붉은색으로 바꿔라" in prompt or "를 붉은색으로 바꿔라" in prompt:
         prompt = "첫 번째 이미지를 붉은색 톤으로 변경해주세요. 전체적인 색상을 붉은 계열로 조정하고 자연스러�� 느낌을 유지해주세요."
     
-    # 명확한 이미지 생성 요청 추가
     prompt += " 이미지를 생성해주세요."
-    
     return prompt
 
-def generate_with_images(prompt, images, max_retries=2):
+def generate_with_images(prompt, images):
     """
-    공식 문서에 기반한 올바른 API 호출 방식 구현
-    실패 시 재시도 기능 추가
+    API 호출을 통해 이미지를 생성하는 함수
     """
-    retries = 0
-    last_error = None
-    
-    while retries <= max_retries:
-        try:
-            # API 키 확인
-            api_key = os.environ.get("GEMINI_API_KEY")
-            if not api_key:
-                return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
-                
-            # Gemini 클라이언트 초기화
-            client = genai.Client(api_key=api_key)
-            
-            logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt} (시도: {retries+1}/{max_retries+1})")
-            
-            # 이미지 및 텍스트 컨텐츠 준비
-            parts = []
-            
-            # 이미지 파트 추가
-            for idx, img in enumerate(images, 1):
-                if img is not None:
-                    # 이미지를 바이트로 변환
-                    buffered = BytesIO()
-                    img.save(buffered, format="PNG")
-                    image_bytes = buffered.getvalue()
-                    
-                    # 이미지 파트 생성
-                    parts.append({
-                        'inlineData': {
-                            'mimeType': 'image/png',
-                            'data': base64.b64encode(image_bytes).decode('utf-8')
-                        }
-                    })
-                    logger.info(f"이미지 #{idx} 추가됨")
-            
-            # 텍스트 프롬프트 추가
-            parts.append({
-                'text': prompt
-            })
+    try:
+        api_key = os.environ.get("GEMINI_API_KEY")
+        if not api_key:
+            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
             
-            # 모델 생성 요청
-            request = {
-                'contents': [{
-                    'parts': parts
-                }],
-                'generation_config': {
-                    'temperature': 1.0,
-                    'topP': 0.95,
-                    'topK': 40,
-                    'maxOutputTokens': 8192,
-                    'responseMimeType': 'image/png'
-                }
-            }
+        client = genai.Client(api_key=api_key)
+        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
+        
+        contents = []
+        contents.append(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
+            )
+        )
+        
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            temp_path = tmp.name
             
-            # 이미지 생성 API 호출
-            model = client.get_model('gemini-pro-vision')
-            response = model.generate_content(request)
+            result_text = ""
+            image_found = False
             
-            # 응답 유효성 확인
-            if not response or not response.parts:
-                if retries < max_retries:
-                    retries += 1
-                    logger.warning(f"유효한 응답을 받지 못했습니다. 재시도 중... ({retries}/{max_retries})")
-                    continue
-                return None, "이미지 생성에 실패했습니다. 유효한 응답을 받지 못했습니다."
+            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("응답에서 이미지 추출 성공")
             
-            # 임시 파일 생성
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-                temp_path = tmp.name
-                
-                result_text = ""
-                image_found = False
-                
-                # 이미지 데이터 확인 및 저장
-                for part in response.parts:
-                    if hasattr(part, 'mime_type') and part.mime_type.startswith('image/'):
-                        # base64 디코딩 및 저장
-                        img_data = base64.b64decode(part.data)
-                        with open(temp_path, 'wb') as f:
-                            f.write(img_data)
-                        image_found = True
-                        logger.info("응답에서 이미지 추출 성공")
-                        break
-                
-                # 텍스트 추출 (옵션)
-                if hasattr(response, 'text'):
-                    result_text = response.text
-                
-                if not image_found:
-                    if retries < max_retries:
-                        retries += 1
-                        logger.warning(f"API에서 이미지를 생성하지 못했습니다. 재시도 중... ({retries}/{max_retries})")
-                        continue
-                    return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
-                        
-                # 결과 이미지 반환
-                result_img = Image.open(temp_path)
-                if result_img.mode == "RGBA":
-                    result_img = result_img.convert("RGB")
+            if not image_found:
+                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
                     
-                return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
+            result_img = Image.open(temp_path)
+            if result_img.mode == "RGBA":
+                result_img = result_img.convert("RGB")
                 
-        except Exception as e:
-            last_error = str(e)
-            logger.exception(f"이미지 생성 중 오류 발생 (시도: {retries+1}/{max_retries+1}):")
-            if retries < max_retries:
-                retries += 1
-                logger.info(f"재시도 중... ({retries}/{max_retries})")
-                continue
-            return None, f"오류 발생: {last_error}"
-    
-    return None, f"최대 재시도 횟수 초과. 마지막 오류: {last_error}"
-
+            return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
+            
+    except Exception as e:
+        logger.exception("이미지 생성 중 오류 발생:")
+        return None, f"오류 발생: {str(e)}"
 
 def process_images_with_prompt(image1, image2, image3, prompt):
     """
-    3개의 이미지와 프롬프트를 처리하는 함수
-    이미지 없이 프롬프트만으로도 동작
+    3개의 이미지와 프롬프트를 처리하는 함수.
+    이미지가 업로드되지 않은 경우에도 프롬프트에 따라 이미지를 생성하며,
+    이미지 생성에 실패하면 최대 3회까지 재시도합니다.
     """
     try:
-        # 이미지 개수 확인
         images = [image1, image2, image3]
-        valid_images = [img for img in images if img is not None]
+        valid_images = [img for img in images if img is not None]  # 이미지가 없을 수도 있음
         
-        # 프롬프트 처리
+        # 프롬프트 자동 생성 (업로드된 이미지 수에 따라)
         if not prompt or not prompt.strip():
-            # 프롬프트가 없는 경우
-            if not valid_images:
-                # 이미지도 없고 프롬프트도 없는 경우 기본 프롬프트 사용
-                prompt = "아름다운 풍경 이미지를 생성해주세요. 푸른 하늘, 초록색 언덕, 그리고 맑은 호수가 있는 평화로운 자연 풍경을 만들어주세요."
-                logger.info("이미지와 프롬프트 모두 없음. 기본 프롬프트 생성")
+            if len(valid_images) == 0:
+                prompt = "이 프롬프트에 맞춰 창의적으로 이미지를 생성해주세요."
+                logger.info("텍스트 전용 프롬프트 자동 생성")
             elif len(valid_images) == 1:
                 prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
                 logger.info("단일 이미지 프롬프트 자동 생성")
@@ -234,21 +167,30 @@ def process_images_with_prompt(image1, image2, image3, prompt):
                 prompt = "이 세 이미지를 창의적으로 합성해주세요. 모든 이미지의 주요 요소를 포함하되 자연스럽고 일관된 하나의 장면으로 만들어주세요."
                 logger.info("세 이미지 합성 프롬프트 자동 생성")
         else:
-            # 프롬프트 전처리 및 기능 명령 해석
             prompt = preprocess_prompt(prompt, image1, image2, image3)
         
-        # 이미지 생성
-        return generate_with_images(prompt, valid_images)
+        # 최대 3회까지 재시도
+        max_retries = 3
+        for attempt in range(max_retries):
+            result_img, status = generate_with_images(prompt, valid_images)
+            if result_img is not None:
+                return result_img, status
+            else:
+                logger.info(f"이미지 생성 실패, 재시도 중 ({attempt+1}/{max_retries})")
+        return None, f"이미지 생성에 실패하였습니다. 마지막 상태: {status}"
             
     except Exception as e:
         logger.exception("이미지 처리 중 오류 발생:")
         return None, f"오류 발생: {str(e)}"
 
-# 기능 선택 콜백
 def update_prompt_from_function(function_choice):
+    """
+    기능 선택에 따라 프롬프트 템플릿을 반환합니다.
+    (커스텀 텍스트 입력은 제거되었습니다.)
+    """
     function_templates = {
-        "1. 이미지 변경": "#1을 창의적으로 바꿔라",
-        "2. 글자지우기": "#1에서 텍스트를 지워라",
+        "1. 이미지 변경": '#1을 "원하는 설명"으로 바꿔라',
+        "2. 글자지우기": '#1에서 "지울 텍스트"를 지워라',
         "3. 얼굴바꾸기": "#1의 인물을 #2의 얼굴로 바꿔라",
         "4. 옷바꾸기": "#1의 인물에 #2 또는 #3의 옷으로 변경하라",
         "5. 배경바꾸기": "#1의 이미지에 #2의 배경으로 자연스럽게 바꿔라",
@@ -258,26 +200,26 @@ def update_prompt_from_function(function_choice):
     
     return function_templates.get(function_choice, "")
 
-# Gradio 인터페이스
+# Gradio 인터페이스 구성
 with gr.Blocks() as demo:
     gr.HTML(
         """
         <div style="text-align: center; margin-bottom: 1rem;">
             <h1>간단한 이미지 생성기</h1>
-            <p>이미지를 업로드하거나 프롬프트를 입력하여 이미지를 생성하세요. 둘 다 없어도 기본 이미지가 생성됩니다.</p>
+            <p>프롬프트만 입력해도 이미지가 생성됩니다. 또한, 기능 선택만으로 프롬프트에 적용할 수 있습니다.</p>
         </div>
         """
     )
     
     with gr.Row():
         with gr.Column():
-            # 3개의 이미지 입력
+            # 3개의 이미지 입력 (업로드하지 않아도 됨)
             with gr.Row():
                 image1_input = gr.Image(type="pil", label="#1", image_mode="RGB")
                 image2_input = gr.Image(type="pil", label="#2", image_mode="RGB")
                 image3_input = gr.Image(type="pil", label="#3", image_mode="RGB")
             
-            # 기능 선택 드롭다운
+            # 기능 선택 드롭다운 (커스텀 텍스트 입력 제거)
             function_dropdown = gr.Dropdown(
                 choices=[
                     "1. 이미지 변경", 
@@ -295,7 +237,7 @@ with gr.Blocks() as demo:
             # 프롬프트 입력 (선택 사항)
             prompt_input = gr.Textbox(
                 lines=3,
-                placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다. 이미지 없이도 프롬프트만으로 생성 가능합니다.",
+                placeholder="프롬프트를 입력하거나 비워두면 자동 생성됩니다.",
                 label="프롬프트 (선택 사항)"
             )
             
@@ -310,7 +252,7 @@ with gr.Blocks() as demo:
             # 사용된 프롬프트 표시
             prompt_display = gr.Textbox(label="사용된 프롬프트", visible=True)
     
-    # 기능 선택 드롭다운 이벤트 (즉시 프롬프트 업데이트)
+    # 기능 선택 드롭다운 변경 시 프롬프트 업데이트 (apply 버튼 없이 바로 적용)
     function_dropdown.change(
         fn=update_prompt_from_function,
         inputs=[function_dropdown],
@@ -324,12 +266,10 @@ with gr.Blocks() as demo:
         valid_images = [img for img in images if img is not None]
         
         try:
-            # 자동 프롬프트 생성 또는 프롬프트 전처리
             auto_prompt = prompt
             if not prompt or not prompt.strip():
-                if not valid_images:
-                    # 이미지도 없고 프롬프트도 없는 경우
-                    auto_prompt = "아름다운 풍경 이미지를 생성해주세요. 푸른 하늘, 초록색 언덕, 그리고 맑은 호수가 있는 평화로운 자연 풍경을 만들어주세요."
+                if len(valid_images) == 0:
+                    auto_prompt = "이 프롬프트에 맞춰 창의적으로 이미지를 생성해주세요."
                 elif len(valid_images) == 1:
                     auto_prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
                 elif len(valid_images) == 2:
@@ -339,7 +279,6 @@ with gr.Blocks() as demo:
             else:
                 auto_prompt = preprocess_prompt(prompt, image1, image2, image3)
             
-            # 이미지 생성 함수 호출
             result_img, status = process_images_with_prompt(image1, image2, image3, prompt)
             
             return result_img, status, auto_prompt
@@ -357,17 +296,13 @@ with gr.Blocks() as demo:
         """
         ### 사용 방법:
         
-        1. **자��� 합성**: 이미지만 업로드하고 프롬프트를 비워두면 자동으로 합성됩니다
-        2. **텍스트만으로 생성**: 이미지 없이 프롬프트만 입력해도 이미지가 생성됩니다
-        3. **기능 사용**: 드롭다운에서 원하는 기능을 선택하면 자동으로 프롬프트가 입력됩니다
-        4. **이미지 참조**: #1, #2, #3으로 각 이미지를 참조할 수 있습니다
-        5. **일부 이미지만**: 필요한 이미지만 업로드해도 기능 실행이 가능합니다
-        6. **재시도 기능**: 이미지 생성 중 오류 발생 시 자동으로 재시도합니다
-        
-        > **팁**: 기능을 선택한 후에도 프롬프트를 직접 수정할 수 있습니다
+        1. **텍스트 전용 생성**: 이미지를 업로드하지 않고 프롬프트만 입력해도 이미지를 생성할 수 있습니다.
+        2. **기능 선택**: 드롭다운에서 원하는 기능을 선택하면 해당 기능의 템플릿이 프롬프트에 자동 반영됩니다.
+        3. **이미지 업로드**: 일부 또는 전체 이미지를 업로드해도 기능이 정상적으로 동작합니다.
+        4. **자동 재시도**: 이미지 생성 중 오류가 발생하면 최대 2회 추가 재시도합니다.
         """
     )
 
 # 애플리케이션 실행
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)