from typing import Optional
import gradio as gr

import quantize
from huggingface_hub import HfApi, login

def run(model_id: str, model_version: str, additional_args: str, token: Optional[str] = None) -> str:
    if model_id == "":
        return "Please enter model_id."
    login(token=token)
    api = HfApi(token=token)

    quantize.quantize(api=api, model_id=model_id, model_version=model_version, additional_args=additional_args)


DESCRIPTION = """
Simple utility tool to quantize diffusion models and convert them to CoreML.
"""

title="Quantize model and convert to CoreML"

with gr.Blocks(title=title) as demo:
    description = gr.Markdown(f"""# {title}""")
    description = gr.Markdown(DESCRIPTION)

    with gr.Row() as r:
        with gr.Column() as c:
            model_id = gr.Text(max_lines=1, label="ID of output repo")
            model_version = gr.Text(max_lines=1, label="Version of model to convert", value="stabilityai/sd-turbo")
            additional_args = gr.Text(max_lines=1, label="Additional Args (optional)")
            token = gr.Text(max_lines=1, label="Your HuggingFace write token")
            with gr.Row() as c:
                clean = gr.ClearButton()
                submit = gr.Button("Submit", variant="primary")

        with gr.Column() as d:
            output = gr.Markdown()

    submit.click(run, inputs=[model_id, model_version, additional_args, token], outputs=output, concurrency_limit=1)

demo.queue(max_size=10).launch(show_api=True)