Update app.py

app.py

@@ -57,29 +57,37 @@ def update_imgbox(choices):
 executor = ThreadPoolExecutor(max_workers=num_models)
 
 
+#def gen_fn(model_str, prompt):
+#    if model_str == 'NA':
+#        return None#
+#
+#    noise = str(randint(0, 9999))
+#    combined_prompt = f'{prompt}'
+#    print(f"Generating with prompt: {combined_prompt}")
+
 def gen_fn(model_str, prompt):
+    global models_load
     if model_str == 'NA':
         return None
-
+    if models_load.get(model_str) is None:
+        models_load[model_str] = load_model(model_str)  # Load the model if not already loaded
     noise = str(randint(0, 9999))
-    combined_prompt = f'{prompt}'
-    print(f"Generating with prompt: {combined_prompt}")
+    return models_load[model_str](f'{prompt} {noise}')
+
     
-    try:
-        image_response = models_load[model_str](f'{prompt} {noise}')
-      # print(f"77 {models_load[model_str](f'{combined_prompt}')}")
-    #   image_response = models_load[model_str](f'{combined_prompt}')
-        # 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
+   # try:
+   #     image_response = models_load[model_str](f'{prompt} {noise}')
+     
+    #    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