Update app.py

app.py

@@ -416,13 +416,12 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             if randomize_seed:
                 seed = random.randint(0, MAX_SEED)
 
+
         # Generate image
         if image_input is not None:
             final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-            return final_image, seed, gr.update(visible=False)
         else:
             image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
-            # Consume the generator to get the final image
             final_image = None
             step_counter = 0
             for image in image_generator:
@@ -431,13 +430,17 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
                 progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
                 yield image, seed, gr.update(value=progress_bar, visible=True)
             
-            if final_image is None:
-                raise gr.Error("Failed to generate image")
-            
-            return final_image, seed, gr.update(value=progress_bar, visible=False)
+        if final_image is None:
+            raise gr.Error("Failed to generate image")
+        
+        return final_image, seed, gr.update(visible=False)
+    except spaces.zero.gradio.HTMLError as e:
+        if "You have exceeded your Pro GPU quota" in str(e):
+            return gr.Error("GPU quota exceeded. Please try again later."), seed, gr.update(visible=False)
+        raise
     except Exception as e:
         print(f"Error in run_lora: {str(e)}")
-        return None, seed, gr.update(visible=False)
+        return gr.Error(f"An error occurred: {str(e)}"), seed, gr.update(visible=False)
 
 run_lora.zerogpu = True
 
@@ -563,8 +566,8 @@ def infer_upscale(
 ):
     try:
         if input_image is None:
-            raise ValueError("No input image provided")
-        
+            return gr.Error("No input image provided"), seed
+
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
@@ -603,7 +606,7 @@ def infer_upscale(
         return image, seed
     except Exception as e:
         print(f"Error in infer_upscale: {str(e)}")
-        return None, seed
+        return gr.Error(f"Upscaling failed: {str(e)}"), seed
 
 def check_upscale_input(input_image, *args):
     if input_image is None:
@@ -755,29 +758,31 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         inputs=[selected_indices, loras_state],
         outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
-    # 이벤트 핸들러 수정
+
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
         inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
         outputs=[result, seed, progress_bar]
     ).then(
-        fn=lambda x, history: update_history(x, history) if x is not None else history,
+        fn=lambda x, history: update_history(x, history) if not isinstance(x, gr.Error) else history,
         inputs=[result, history_gallery],
         outputs=history_gallery,
     )
-
-    # 업스케일 버튼 이벤트 추가 (이 부분을 gr.Blocks 컨텍스트 내부로 이동)
+    
     upscale_button.click(
-        check_upscale_input,
+        infer_upscale,
         inputs=[
-            upscale_input,
             upscale_seed,
             upscale_randomize_seed,
+            upscale_input,
             upscale_steps,
             upscale_factor,
             controlnet_conditioning_scale,
         ],
+        outputs=[upscale_result, upscale_seed_output],
+    )
+    
         outputs=[
             upscale_input,
             upscale_seed,