Update app.py

app.py

@@ -25,7 +25,7 @@ article_text = """
 
 def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024, output_format="PNG"):
     if prompt == "" or prompt == None:
-        return None
+        return None, None, None
 
     if lora_id.strip() == "" or lora_id == None:
         lora_id = "black-forest-labs/FLUX.1-dev"
@@ -71,10 +71,11 @@ def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1
         img_byte_arr = io.BytesIO()
         image.save(img_byte_arr, format=output_format)
         img_byte_arr.seek(0)
-        return Image.open(img_byte_arr), seed, seed, img_byte_arr.getvalue()
+        
+        return Image.open(io.BytesIO(img_byte_arr.getvalue())), seed, io.BytesIO(img_byte_arr.getvalue())
     except Exception as e:
         print(f"Error when trying to open the image: {e}")
-        return None
+        return None, None, None
 
 
 examples = [
@@ -119,7 +120,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         with gr.Row():
             seed_output = gr.Textbox(label="Seed Used", show_copy_button=True, elem_id="seed-output")
         with gr.Row():
-            download_button = gr.File(label="Download Image", file_name="generated_image", elem_id="download-button")
+            download_button = gr.File(label="Download Image", elem_id="download-button")
 
         gr.Markdown(article_text)
         
@@ -128,12 +129,10 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
             inputs=[text_prompt],
         )
 
-        def download_image(image, output_format):
-            format_extension = output_format.lower()
-            file_name = f"generated_image.{format_extension}"
-            return file_name, image
+        def download_image(image_data):
+            return image_data
 
-        text_button.click(query, inputs=[custom_lora, text_prompt, steps, cfg, randomize_seed, seed, width, height, output_format], outputs=[image_output, seed_output, seed, download_button])
-        download_button.change(download_image, inputs=[image_output, output_format], outputs=download_button)
+        text_button.click(query, inputs=[custom_lora, text_prompt, steps, cfg, randomize_seed, seed, width, height, output_format], outputs=[image_output, seed_output, download_button])
+        download_button.change(download_image, inputs=[download_button], outputs=download_button)
 
 app.launch(show_api=False, share=True)