Update app.py

app.py

@@ -33,11 +33,23 @@ def update_imgbox(choices):
     choices_plus = extend_choices(choices)
     return [gr.Image(None, label=m, visible=(m != 'NA'), elem_id="custom_image") for m in choices_plus]
 
-def gen_fn(model_str, prompt):
+def gen_fn(model_str, prompt, max_retries=10):
     if model_str == 'NA':
         return None
-    noise = str(randint(0, 9999999))
-    return models_load[model_str](f'{prompt} {noise}')
+    
+    retries = 0
+    while retries < max_retries:
+        try:
+            noise = str(randint(0, 9999999))
+            result = models_load[model_str](f'{prompt} {noise}')
+            return result
+        except Exception as e:
+            print(f"Error generating image: {e}")
+            retries += 1
+            if retries >= max_retries:
+                raise Exception(f"Failed to generate image after {max_retries} retries.")
+    
+    return None
 
 def make_me():
     with gr.Row():