Update app.py

app.py

@@ -213,39 +213,55 @@ def generate_with_images(prompt, images, variation_index=0):
         logger.exception("이미지 생성 중 오류 발생:")
         return None, f"오류 발생: {str(e)}"
 
-def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0):
+def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0, max_retries=3):
     """
     3개의 이미지와 프롬프트를 처리하여 최종 영어 프롬프트(final_prompt)를 생성한 후,
-    API를 호출하여 결과 이미지를 반환합니다.
+    API를 호출하여 결과 이미지를 반환합니다. 에러 발생 시 최대 max_retries 횟수만큼 재시도합니다.
     """
-    try:
-        images = [image1, image2, image3]
-        valid_images = [img for img in images if img is not None]
-        if not valid_images:
-            return None, "적어도 하나의 이미지를 업로드해주세요.", ""
+    retry_count = 0
+    last_error = None
+    
+    while retry_count < max_retries:
+        try:
+            images = [image1, image2, image3]
+            valid_images = [img for img in images if img is not None]
+            if not valid_images:
+                return None, "적어도 하나의 이미지를 업로드해주세요.", ""
 
-        if prompt and prompt.strip():
-            processed_prompt = preprocess_prompt(prompt, image1, image2, image3)
-            if re.search("[가-힣]", processed_prompt):
-                final_prompt = translate_prompt_to_english(processed_prompt)
+            if prompt and prompt.strip():
+                processed_prompt = preprocess_prompt(prompt, image1, image2, image3)
+                if re.search("[가-힣]", processed_prompt):
+                    final_prompt = translate_prompt_to_english(processed_prompt)
+                else:
+                    final_prompt = processed_prompt
             else:
-                final_prompt = processed_prompt
-        else:
-            if len(valid_images) == 1:
-                final_prompt = "Please creatively transform this image into a more vivid and artistic version. Do not include any text or watermarks in the generated image."
-                logger.info("Default prompt generated for single image")
-            elif len(valid_images) == 2:
-                final_prompt = "Please seamlessly composite these two images, integrating their key elements harmoniously into a single image. Do not include any text or watermarks in the generated image."
-                logger.info("Default prompt generated for two images")
-            else:
-                final_prompt = "Please creatively composite these three images, combining their main elements into a cohesive and natural scene. Do not include any text or watermarks in the generated image."
-                logger.info("Default prompt generated for three images")
+                if len(valid_images) == 1:
+                    final_prompt = "Please creatively transform this image into a more vivid and artistic version. Do not include any text or watermarks in the generated image."
+                    logger.info("Default prompt generated for single image")
+                elif len(valid_images) == 2:
+                    final_prompt = "Please seamlessly composite these two images, integrating their key elements harmoniously into a single image. Do not include any text or watermarks in the generated image."
+                    logger.info("Default prompt generated for two images")
+                else:
+                    final_prompt = "Please creatively composite these three images, combining their main elements into a cohesive and natural scene. Do not include any text or watermarks in the generated image."
+                    logger.info("Default prompt generated for three images")
 
-        result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
-        return result_img, status, final_prompt
-    except Exception as e:
-        logger.exception("이미지 처리 중 오류 발생:")
-        return None, f"오류 발생: {str(e)}", prompt
+            result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
+            if result_img is not None:
+                return result_img, status, final_prompt
+            else:
+                # API 호출은 성공했지만 이미지 생성 실패
+                last_error = status
+                retry_count += 1
+                logger.warning(f"이미지 생성 실패, 재시도 {retry_count}/{max_retries}: {status}")
+                time.sleep(1)  # 1초 대기 후 재시도
+        except Exception as e:
+            last_error = str(e)
+            retry_count += 1
+            logger.exception(f"이미지 처리 중 오류 발생, 재시도 {retry_count}/{max_retries}:")
+            time.sleep(1)  # 1초 대기 후 재시도
+    
+    # 모든 재시도 실패 후
+    return None, f"최대 재시도 횟수({max_retries}회) 초과 후 실패: {last_error}", prompt
 
 def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progress()):
     """
@@ -256,12 +272,13 @@ def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progres
     prompts = []
     
     num_images = 4  # 생성할 이미지 수
+    max_retries = 3  # 각 이미지 당 최대 재시도 횟수
     
     progress(0, desc="이미지 생성 준비 중...")
     
     for i in range(num_images):
         progress((i / num_images), desc=f"{i+1}/{num_images} 이미지 생성 중...")
-        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, i)
+        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, i, max_retries)
         
         if result_img is not None:
             results.append(result_img)
@@ -319,6 +336,7 @@ with gr.Blocks() as demo:
                 text_change_btn = gr.Button("글자변경하기")
                 clothes_change_btn1 = gr.Button("가상 상품착용-1")
                 clothes_change_btn2 = gr.Button("상품착용-2")
+                holding_product_btn = gr.Button("상품들고 있기")
                 background_change_btn = gr.Button("배경바꾸기")
                 composite_product_btn = gr.Button("부분지우기")
             submit_btn = gr.Button("이미지 생성 (4장)", variant="primary")
@@ -380,6 +398,11 @@ with gr.Blocks() as demo:
         inputs=[],
         outputs=prompt_input
     )
+    holding_product_btn.click(
+        fn=lambda: "(#1의 여성모델)이 신체 비율과 포즈는 유치한 체 (#2의 크지 않는 와인잔)를 직접 모델이 사용하기 위해 들고 있는 자연스러운 모습으로 이미지를 생성하라.",
+        inputs=[],
+        outputs=prompt_input
+    )
     background_change_btn.click(
         fn=lambda: "(#1의 여성모델)이 (#2 이미지의 배경)을 주요 피사체는 그대로 유지하여 두 이미지의 분위기가 자연스럽게 어우러지도록 생성하라.",
         inputs=[],