Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from random import randint
 from all_models import models
+import os
+from PIL import Image
 
 def load_fn(models):
     global models_load
@@ -32,7 +34,14 @@ def gen_fn(model_str, prompt, negative_prompt=None, image_style="Default"):
     modified_prompt = f"{prompt}, {image_style}"
     if negative_prompt:
         modified_prompt += f", not {negative_prompt}"
-    return models_load[model_str](modified_prompt)
+    image = models_load[model_str](modified_prompt)
+    if image:
+        image_path = f"output/{model_str}_{randint(0, 10000)}.jpg"
+        if not os.path.exists('output'):
+            os.makedirs('output')
+        image.save(image_path, format="JPEG")
+        return image_path
+    return None
 
 with gr.Blocks() as demo:
     with gr.Tab('The Dream'):
@@ -81,4 +90,4 @@ with gr.Blocks() as demo:
             """
             )
 
-demo.launch(max_threads=200)
+demo.launch(max_threads=200)