Update app.py

app.py

@@ -248,12 +248,29 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             }
         }
 
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=100)
-        response.raise_for_status()
+        # 타임아웃 값을 늘리고 재시도 로직 추가
+        max_retries = 3
+        current_retry = 0
         
-        image = Image.open(io.BytesIO(response.content))
-        print(f'Generation {key} completed successfully')
-        return image
+        while current_retry < max_retries:
+            try:
+                response = requests.post(API_URL, headers=headers, json=payload, timeout=180)  # 타임아웃을 180초로 증가
+                response.raise_for_status()
+                
+                image = Image.open(io.BytesIO(response.content))
+                print(f'Generation {key} completed successfully')
+                return image
+                
+            except requests.exceptions.Timeout:
+                current_retry += 1
+                if current_retry < max_retries:
+                    print(f"Timeout occurred. Retrying... (Attempt {current_retry + 1}/{max_retries})")
+                    continue
+                else:
+                    raise gr.Error(f"Request timed out after {max_retries} attempts. The model might be busy, please try again later.")
+                    
+            except requests.exceptions.RequestException as e:
+                raise gr.Error(f"Request failed: {str(e)}")
         
     except requests.exceptions.RequestException as e:
         error_message = f"Request failed: {str(e)}"
@@ -269,8 +286,6 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
         raise gr.Error(f"Unexpected error: {str(e)}")
 
 
-
-
 def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height, progress=gr.Progress()):
     if len(selected_models) > 4:
         raise gr.Error("Please select up to 4 models")