import gradio as gr

from min_dalle import MinDalle
import torch

model = MinDalle(is_mega=True, models_root='./pretrained')

def text_to_image(text, grid_size=1):
    with torch.no_grad():
        return model.generate_image(text, grid_size=grid_size)


iface = gr.Interface(fn=text_to_image, 
                    inputs=[gr.Textbox(),gr.Number(value=1,precision=0)],
                    outputs='image',
                    title='Min-Dalle',
                    description="AI model generating images from any prompt!"
)
iface.launch()