Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 import io
 import random
 import os
+import time
 from PIL import Image
 import json
 
@@ -248,33 +249,45 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             }
         }
 
-        # 타임아웃 값을 늘리고 재시도 로직 추가
+        # Improved retry logic with exponential backoff
         max_retries = 3
         current_retry = 0
+        backoff_factor = 2  # Exponential backoff
         
         while current_retry < max_retries:
             try:
-                response = requests.post(API_URL, headers=headers, json=payload, timeout=180)  # 타임아웃을 180초로 증가
+                response = requests.post(API_URL, headers=headers, json=payload, timeout=180)  # 3-minute timeout
                 response.raise_for_status()
                 
                 image = Image.open(io.BytesIO(response.content))
                 print(f'Generation {key} completed successfully')
                 return image
                 
-            except requests.exceptions.Timeout:
+            except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, 
+                    requests.exceptions.HTTPError, requests.exceptions.RequestException) as e:
                 current_retry += 1
                 if current_retry < max_retries:
-                    print(f"Timeout occurred. Retrying... (Attempt {current_retry + 1}/{max_retries})")
+                    wait_time = backoff_factor ** current_retry  # Exponential backoff
+                    print(f"Network error occurred: {str(e)}. Retrying in {wait_time} seconds... (Attempt {current_retry + 1}/{max_retries})")
+                    time.sleep(wait_time)  # Add delay before retry
                     continue
                 else:
-                    raise gr.Error(f"Request timed out after {max_retries} attempts. The model might be busy, please try again later.")
+                    # Detailed error message based on exception type
+                    if isinstance(e, requests.exceptions.Timeout):
+                        error_msg = f"Request timed out after {max_retries} attempts. The model might be busy, please try again later."
+                    elif isinstance(e, requests.exceptions.ConnectionError):
+                        error_msg = f"Connection error after {max_retries} attempts. Please check your network connection."
+                    elif isinstance(e, requests.exceptions.HTTPError):
+                        status_code = e.response.status_code if hasattr(e, 'response') and e.response is not None else "unknown"
+                        error_msg = f"HTTP error (status code: {status_code}) after {max_retries} attempts."
+                    else:
+                        error_msg = f"Request failed after {max_retries} attempts: {str(e)}"
                     
-            except requests.exceptions.RequestException as e:
-                raise gr.Error(f"Request failed: {str(e)}")
+                    raise gr.Error(error_msg)
         
-    except requests.exceptions.RequestException as e:
-        error_message = f"Request failed: {str(e)}"
-        if hasattr(e, 'response') and e.response is not None:
+    except Exception as e:
+        error_message = f"Unexpected error: {str(e)}"
+        if isinstance(e, requests.exceptions.RequestException) and hasattr(e, 'response') and e.response is not None:
             if e.response.status_code == 401:
                 error_message = "Invalid API token. Please check your Hugging Face API token."
             elif e.response.status_code == 403:
@@ -282,8 +295,6 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             elif e.response.status_code == 503:
                 error_message = "Model is currently loading. Please try again in a few moments."
         raise gr.Error(error_message)
-    except Exception as e:
-        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()):
@@ -292,37 +303,61 @@ def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps,
     if len(selected_models) == 0:
         raise gr.Error("Please select at least 1 model")
     
-    # 초기 이미지 배열 생성
+    # Initialize image array
     images = [None] * 4
     total_models = len(selected_models[:4])
     
     def update_gallery():
-        # None이 아닌 이미지만 포함하여 갤러리 업데이트
+        # Only include non-None images for gallery update
         return [img for img in images if img is not None]
     
-    # 각 모델별로 이미지 생성
+    # Create placeholder for failed models
+    placeholder_image = None
+    
+    # Generate image for each model
     for idx, model_name in enumerate(selected_models[:4]):
         try:
             progress((idx + 1) / total_models, f"Generating image for {model_name}...")
             img = query(prompt, model_name, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height)
             images[idx] = img
-            # 이미지가 생성될 때마다 갤러리 업데이트
+            
+            # If this is the first successful generation, save as placeholder for failed models
+            if placeholder_image is None:
+                placeholder_image = img
+                
+            # Update gallery after each successful generation
             yield update_gallery()
         except Exception as e:
             print(f"Error generating image for {model_name}: {str(e)}")
+            # Keep the slot as None and continue with next model
             continue
     
-    # 남은 슬롯을 마지막 생성된 이미지로 채움
-    last_valid_image = next((img for img in reversed(images) if img is not None), None)
-    if last_valid_image:
+    # Fill empty slots with a placeholder (either the last successful image or a blank image)
+    if placeholder_image:
         for i in range(len(images)):
             if images[i] is None:
-                images[i] = last_valid_image
+                # Create a copy of placeholder to avoid reference issues
+                images[i] = placeholder_image.copy() 
+    else:
+        # If all models failed, create a blank image with error text
+        for i in range(len(images)):
+            blank_img = Image.new('RGB', (width, height), color=(240, 240, 240))
+            images[i] = blank_img
     
     progress(1.0, "Generation complete!")
     yield update_gallery()
 
 
+def check_network_connectivity():
+    """Utility function to check network connectivity to the Hugging Face API"""
+    try:
+        response = requests.get("https://api-inference.huggingface.co", timeout=5)
+        if response.status_code == 200:
+            return True
+        return False
+    except:
+        return False
+
 
 css = """
 footer {
@@ -374,17 +409,17 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     lines=1
                 )
 
-                # 상위 4개 모델을 기본으로 설정
+                # Set top 4 models as default
                 default_models = [
-                    "FLUX.1 [Schnell]",  # 모델 이름 통일
+                    "FLUX.1 [Schnell]",
                     "Stable Diffusion 3.5 Large",
                     "Stable Diffusion 3.5 Large Turbo",
                     "Midjourney"
                 ]
 
-                # 전체 모델 리스트
+                # Full model list
                 models_list = [
-                    "FLUX.1 [Schnell]",  # 모델 이름 통일
+                    "FLUX.1 [Schnell]",
                     "Stable Diffusion 3.5 Large",
                     "Stable Diffusion 3.5 Large Turbo",
                     "Stable Diffusion XL",
@@ -428,7 +463,15 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
     with gr.Row():
         generate_btn = gr.Button("Generate 2x2 Grid", variant="primary", size="lg")
         
-
+    # Add network status indicator
+    network_status = gr.Markdown("", elem_id="network_status")
+    
+    # Function to check and update network status
+    def update_network_status():
+        if check_network_connectivity():
+            return "✅ Connected to Hugging Face API"
+        else:
+            return "❌ No connection to Hugging Face API. Please check your network."
 
     with gr.Row():
         gallery = gr.Gallery(
@@ -438,10 +481,10 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
             columns=2,
             rows=2,
             height="auto",
-            preview=True,  # 실시간 프리뷰 활성화
+            preview=True,
         )
 
-# 이벤트 핸들러 수정
+    # Event handlers
     generate_btn.click(
         fn=generate_grid,
         inputs=[
@@ -465,6 +508,9 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
         return gr.update(choices=filtered_models, value=[])
 
     model_search.change(filter_models, inputs=model_search, outputs=model)
+    
+    # Update network status when the app loads
+    dalle.load(fn=update_network_status, outputs=network_status)
 
 if __name__ == "__main__":
     dalle.launch(show_api=False, share=False)