Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from random import randint
 from all_models import models
+from PIL import Image
+import io
 
 def load_fn(models):
     global models_load
@@ -32,7 +34,16 @@ 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)
+    
+    # Save image as .jpg
+    if image is not None:
+        img = Image.fromarray(image)
+        img_byte_arr = io.BytesIO()
+        img.save(img_byte_arr, format='JPEG')
+        img_byte_arr = img_byte_arr.getvalue()
+        return img_byte_arr
+    return None
 
 with gr.Blocks() as demo:
     with gr.Tab('The Dream'):