trying to fix download button

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from dotenv import load_dotenv
 from transformers import pipeline
+import urllib.request
+from PIL import Image
 
 # Load environment variables from the .env file de forma local
 load_dotenv()
@@ -31,7 +33,14 @@ def get_image(text: str, translate: bool):
         prompt=text, n=1, size="512x512"
     )  # ,response_format="b64_json"
     print(response)
-    return response["data"][0]["url"]  # type: ignore
+
+    urllib.request.urlretrieve(response["data"][0]["url"], "img.png")  # type: ignore
+
+    img = Image.open("img.png")
+
+    return img
+
+    # return response["data"][0]["url"]  # type: ignore
 
 
 with gr.Blocks() as demo:

test.py

@@ -0,0 +1,29 @@
+import gradio as gr
+
+
+def generate_image(input_text):
+    """Generates an image from the input text."""
+
+    # Generate the image
+    image = ...
+
+    # Convert the image to a base64 string
+    base64_string = base64.b64encode(image.tobytes())
+
+    # Return the base64 string
+    return base64_string
+
+
+def download_image(base64_string):
+    """Downloads the output image as a .png file."""
+
+    filename = "image.png"
+    with open(filename, "wb") as f:
+        f.write(base64.b64decode(base64_string))
+
+    # Allow the user to download the file
+    gr.File.download(filename)
+
+
+if __name__ == "__main__":
+    gr.Interface(fn=generate_image, inputs="text", outputs="image").launch()