Update app.py

app.py

@@ -27,6 +27,9 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'Generation {key}: {prompt}')
 
+    try:
+        response = None  # Initialize response variable
+
     # Set API URL based on model selection
     if custom_lora.strip():
         API_URL = f"https://api-inference.huggingface.co/models/{custom_lora.strip()}"
@@ -231,21 +234,20 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
         if model == 'epiCPhotoGasm':
             API_URL = "https://api-inference.huggingface.co/models/Yntec/epiCPhotoGasm"
 
-    # Prepare payload
-    payload = {
-        "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed if seed != -1 else random.randint(1, 1000000000),
-        "strength": strength,
-        "parameters": {
-            "width": width,
-            "height": height
+        # Prepare payload
+        payload = {
+            "inputs": prompt,
+            "is_negative": is_negative,
+            "steps": steps,
+            "cfg_scale": cfg_scale,
+            "seed": seed if seed != -1 else random.randint(1, 1000000000),
+            "strength": strength,
+            "parameters": {
+                "width": width,
+                "height": height
+            }
         }
-    }
 
-    try:
         response = requests.post(API_URL, headers=headers, json=payload, timeout=100)
         response.raise_for_status()
         
@@ -255,13 +257,16 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
         
     except requests.exceptions.RequestException as e:
         error_message = f"Request failed: {str(e)}"
-        if response.status_code == 401:
-            error_message = "Invalid API token. Please check your Hugging Face API token."
-        elif response.status_code == 403:
-            error_message = "Access denied. Please check your API token permissions."
-        elif response.status_code == 503:
-            error_message = "Model is currently loading. Please try again in a few moments."
+        if response:
+            if response.status_code == 401:
+                error_message = "Invalid API token. Please check your Hugging Face API token."
+            elif response.status_code == 403:
+                error_message = "Access denied. Please check your API token permissions."
+            elif 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()):
     if len(selected_models) > 4:
@@ -270,7 +275,7 @@ def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps,
         raise gr.Error("Please select at least 1 model")
     
     # 초기 이미지 배열 생성
-    images = [None] * 4
+    images = []
     total_models = len(selected_models[:4])
     
     # 각 모델별로 이미지 생성
@@ -278,20 +283,37 @@ def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps,
         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 img:
+                images.append(img)
         except Exception as e:
             print(f"Error generating image for {model_name}: {str(e)}")
             continue
     
-    # 남은 슬롯을 마지막 생성된 이미지로 채움
-    last_valid_image = next((img for img in reversed(images) if img is not None), None)
-    if last_valid_image:
-        for i in range(len(images)):
-            if images[i] is None:
-                images[i] = last_valid_image
+    # 최소한 하나의 이미지가 생성되었는지 확인
+    if not images:
+        raise gr.Error("Failed to generate any images. Please try again.")
+    
+    # 4개의 이미지 슬롯을 채움
+    while len(images) < 4:
+        images.append(images[-1])
+    
+    # 이미지가 올바르게 로드되었는지 확인
+    valid_images = []
+    for img in images:
+        try:
+            # 이미지 복사본 생성
+            img_copy = img.copy()
+            valid_images.append(img_copy)
+        except Exception as e:
+            print(f"Error processing image: {str(e)}")
+            # 오류가 발생한 경우 마지막 유효한 이미지로 대체
+            if valid_images:
+                valid_images.append(valid_images[-1].copy())
+            else:
+                raise gr.Error("Failed to process images. Please try again.")
     
     progress(1.0, "Generation complete!")
-    return images
+    return valid_images
 
 css = """
 footer {