# https://huggingface.co/yuntian-deng/latex2im_ss_finetunegptneo
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')

title = "<h1><center>Markup-to-Image Diffusion Models with Scheduled Sampling</center></h1>"
authors = "<center>Yuntian Deng, Noriyuki Kojima, Alexander M. Rush</center>"
info = '<center><a href="https://arxiv.org/pdf/2210.05147.pdf">Paper</a> <a href="https://github.com/da03/markup2im">Code</a></center>'

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(authors)
    gr.Markdown(info)
    with gr.Row():
        with gr.Column(scale=2):
            textbox = gr.Textbox(label=r'Type LaTeX formula below and click "Generate"', lines=1, max_lines=1, placeholder='Type LaTeX formula here and click "Generate"', value=r'\sum_{t=1}^T\E_{y_t \sim {\tilde P(y_t| y_0)}} \left\| \frac{y_t - \sqrt{\bar{\alpha}_t}y_0}{\sqrt{1-\bar{\alpha}_t}} - \epsilon_\theta(y_t, t)\right\|^2.')
            submit_btn = gr.Button("Generate", elem_id="btn")
        with gr.Column(scale=3):
            slider = gr.Slider(0, 1000, value=0, label='step (out of 1000)')
            image = gr.Image(label="Rendered Image", show_label=False, elem_id="image")
    inputs = [textbox]
    outputs = [slider, image, submit_btn]
    def infer(formula):
        data = {'formula': formula, 'api_key': API_KEY}
        try:
            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, submit_btn.update(visible=False)
                yield i, q, submit_btn.update(visible=True)
        except Exception as e:
            yield 1000, 255*np.ones((64, 320, 3)), submit_btn.update(visible=True)
    submit_btn.click(fn=infer, inputs=inputs, outputs=outputs)
demo.queue(concurrency_count=20, max_size=200).launch(enable_queue=True)