test.py

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()