import gradio as gr
import spaces  # ZERO GPU

from tagger import (
    convert_tags_to_ja,
    translate_prompt_to_ja,
)


DESCRIPTION_MD = """
# Convert SD tags to Japanese? tags
- I have no idea if it will be of any use.
""".strip()

DESCRIPTION_MD2 = """
The dictionary was generated using the following repository: [ponapon280/danbooru-e621-converter](https://github.com/ponapon280/danbooru-e621-converter)
""".strip()


def demo():
    with gr.Blocks() as ui:
        gr.Markdown(DESCRIPTION_MD)
        
        with gr.Row():
            with gr.Column():
                with gr.Group():
                    input_general = gr.TextArea(
                        label="Input tags",
                        lines=6,
                        placeholder="1girl, solo, ...",
                    )
                    start_btn = gr.Button(value="CONVERT", size="lg", variant="primary")

            with gr.Column():
                with gr.Group():
                    output_prompt = gr.TextArea(label="Output tags", lines=6, show_copy_button=True, interactive=False)

        gr.Markdown(DESCRIPTION_MD2)

        start_btn.click(convert_tags_to_ja, inputs=[input_general], outputs=[output_prompt]).then(
            translate_prompt_to_ja, inputs=[output_prompt], outputs=[output_prompt],
        )
    
    return ui

if __name__ == "__main__":
    demo().queue().launch()