Update app.py

app.py

@@ -21,13 +21,34 @@ def load_fn(models):
     models_load = {}
     for model in models:
         if model not in models_load.keys():
+          # try:
+          #      m = gr.load(f'models/{model}')
+          #  except Exception as error:
+          #      print(f"Error loading model {model}: {error}")
+          #      m = gr.Interface(lambda _: None, inputs=gr.Textbox(), outputs=gr.Image(), enable_queue=False)
+          #  models_load.update({model: m})
             try:
-                m = gr.load(f'models/{model}')
-            except Exception as error:
-                print(f"Error loading model {model}: {error}")
-                m = gr.Interface(lambda _: None, inputs=gr.Textbox(), outputs=gr.Image(), enable_queue=False)
-            models_load.update({model: m})
-    
+        # Attempt to generate the image
+        image_response = models_load[model_str](f'{prompt} {negative_prompt} {noise}')
+        
+        # Check if the image_response is a tuple, handle accordingly
+        if isinstance(image_response, tuple):
+            # If the response is a tuple, return the first item assuming it's the image
+            image_response = image_response[0]
+        
+        # Ensure the response is an image or image-like object
+        if isinstance(image_response, gr.Image):
+            return image_response
+        elif isinstance(image_response, str):  # If the response is a path or URL, pass it as a string
+            return gr.Image(image_response)  # You can handle it based on your model's return type
+        else:
+            print(f"Unexpected response type: {type(image_response)}")
+            return None
+            
+    except Exception as e:
+        print(f"Error occurred: {e}")
+        return None
+
 
 load_fn(models)