Update app.py

app.py

@@ -28,6 +28,14 @@ def update_imgbox(choices):
     choices_plus = extend_choices(choices)
     return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_plus]
 
+def save_image_as_jpg(image):
+    if image:
+        image = Image.open(io.BytesIO(image))
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='JPEG')
+        return img_byte_arr.getvalue()
+    return None
+
 def gen_fn(model_str, prompt, negative_prompt=None, image_style="Default"):
     if model_str == 'NA':
         return None
@@ -35,15 +43,7 @@ def gen_fn(model_str, prompt, negative_prompt=None, image_style="Default"):
     if negative_prompt:
         modified_prompt += f", not {negative_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
+    return save_image_as_jpg(image)
 
 with gr.Blocks() as demo:
     with gr.Tab('The Dream'):