Update app.py

app.py

@@ -1,42 +1,49 @@
 import os
-os.system("pip install flask")
+
+# Install Flask if not already installed
+return_code = os.system('pip install flask')
+if return_code != 0:
+    raise RuntimeError("Failed to install Flask")
+
+import gradio as gr
 from random import randint
+from all_models import models
 from flask import Flask, request, send_file
 from io import BytesIO
 from PIL import Image
-import importlib
 
 app = Flask(__name__)
 
-def load_model(model_name):
-    try:
-        # Dynamically import the model module
-        module = importlib.import_module(f'all_models.{model_name}')
-        return module.get_model()  # Assume get_model() returns a callable model
-    except Exception as error:
-        print(f"Error loading model {model_name}: {error}")
-        return lambda txt: None  # Return a placeholder function
+def load_fn(models):
+    global models_load
+    models_load = {}
+    
+    for model in models:
+        if model not in models_load.keys():
+            try:
+                m = gr.load(f'models/{model}')
+            except Exception as error:
+                m = gr.Interface(lambda txt: None, ['text'], ['image'])
+            models_load.update({model: m})
+
+load_fn(models)
 
 num_models = 6
-default_models = ['model1', 'model2', 'model3', 'model4', 'model5', 'model6']  # Replace with actual model names
+default_models = models[:num_models]
 
 def extend_choices(choices):
     return choices + (num_models - len(choices)) * ['NA']
 
 def update_imgbox(choices):
     choices_plus = extend_choices(choices)
-    return [Image.open(None) if m != 'NA' else None for m in choices_plus]
+    return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_plus]
 
-def gen_fn(model_str, prompt, negative_prompt=None):
+def gen_fn(model_str, prompt):
     if model_str == 'NA':
         return None
     noise = str('')  # Optional: str(randint(0, 99999999999))
     try:
-        model = load_model(model_str)
-        full_prompt = f'{prompt} {noise}'
-        if negative_prompt:
-            full_prompt += f' -{negative_prompt}'
-        result = model(full_prompt)
+        result = models_load[model_str](f'{prompt} {noise}')
         # Check if result is an image or a file path
         if isinstance(result, str):  # Assuming result might be a file path
             if os.path.exists(result):
@@ -57,14 +64,13 @@ def gen_fn(model_str, prompt, negative_prompt=None):
 def home():
     prompt = request.args.get('prompt', '')
     model = request.args.get('model', default_models[0] if default_models else 'NA')
-    negative_prompt = request.args.get('Nprompt', None)
     
-    if not model:
+    if not model or model not in models_load:
         return f'Invalid model: {model}', 400
     
     if prompt:
         # Generate the image
-        image = gen_fn(model, prompt, negative_prompt)
+        image = gen_fn(model, prompt)
         if isinstance(image, Image.Image):  # Ensure the result is a PIL image
             # Save image to BytesIO object
             img_io = BytesIO()
@@ -76,4 +82,4 @@ def home():
 
 if __name__ == '__main__':
     # Launch Flask app
-    app.run(host='0.0.0.0', port=7860)  # Run Flask app
+    app.run(host='0.0.0.0', port=7860)  # Run Flask app