import gradio as gr
import numpy as np
import requests
import base64
import os

API_ENDPOINT = os.getenv('API_ENDPOINT')
API_KEY = os.getenv('API_KEY')

# setup
gallery = gr.Gallery(label="Rendered Image", show_label=False, elem_id="gallery").style(grid=[1], height="auto")

# infer
def infer(latex): 
    formula = latex
    data = {'formula': formula, 'api_key': API_KEY}
    with requests.post(url=API_ENDPOINT, data=data, timeout=600, stream=True) as r:
        i = 0
        for line in r.iter_lines():
            response = line.decode('ascii').strip()
            r = base64.decodebytes(response.encode('ascii'))
            q = np.frombuffer(r, dtype=np.float32).reshape((64, 320, 3))
            i += 1
            yield i, [q,]
    
title = "Markup-to-Image Diffusion Models with Scheduled Sampling"
description="Yuntian Deng, Noriyuki Kojima, Alexander M. Rush"

# launch
gr.Interface(fn=infer, inputs=["text"], outputs=[gr.Slider(0, 1000, value=0, label='step (out of 1000)'), gallery],title=title,description=description).queue(max_size=100).launch(enable_queue=True)