from transformers import pipeline, set_seed
import gradio as grad, random, re
import os
import sys


gpt2_pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', tokenizer='gpt2')


def generate(starting_text):
    with open("ideas.txt", "r") as f:
        line = f.readlines()
    seed = random.randint(100, 1000000)
    set_seed(seed)

    if starting_text == "":
        starting_text: str = line[random.randrange(0, len(line))].replace("\n", "").capitalize()
        starting_text: str = re.sub(r"\.", '', starting_text)

    response = gpt2_pipe(starting_text, max_length=(len(starting_text) + random.randint(60, 80)), num_return_sequences=1)
    response_list = []
    for x in response:
        resp = x['generated_text'].strip()
        if resp != starting_text and len(resp) > (len(starting_text) + 4) and resp.endswith((":", "-", "—")) is False:
            response_list.append(resp)

    response_end = "\n".join(response_list)
    response_end = re.sub('[^ ]+\.[^ ]+','', response_end)
    response_end = response_end.replace("<", "").replace(">", "")

    if response_end != "":
        return response_end

with grad.Blocks(css='style.css') as demo:
    grad.HTML(
        """
            <div style="text-align: center; max-width: 650px; margin: 0 auto;">
              <div>
                <h2 style="font-weight: 900; font-size: 3rem; margin-bottom:20px;">
                  SD指令生成器.
                </h2>
              </div>
              <p style="margin-bottom: 10px; font-size: 96%">
                输入一个英文单词，然后戳【咒语显现】按钮~ 不知道输入什么？直接戳！
              </p>
              <p style="margin-bottom: 10px; font-size: 98%">
              ❤️ 后台有上千个单词随机组合喔~ ❤️</a>
              </p>
            </div>
        """
    )
    with grad.Column(elem_id="col-container"):
        with grad.Row(variant="compact"):
            txt = grad.Textbox(
                label="Initial Text",
                show_label=False,
                max_lines=1,
                placeholder="输入一两个英语单词",
            ).style(
                container=False,
            )
            run = grad.Button("✨ 咒语显现 ✨").style(full_width=False)



        with grad.Row(variant="compact"):
            out = grad.Textbox(
                label="Generated Text",
                show_label=False,
                lines=5,
            ).style(
                container=False,
            )

        run.click(generate, inputs=[txt], outputs=[out])
        


        with grad.Row():
                grad.HTML(
    """
        <div class="footer">
        <p> Powered by <a href="https://huggingface.co/Gustavosta">Gustavosta</a> Stable Diffusion model
</p>
</div>
        <div class="acknowledgments" style="font-size: 115%">
            <p> <a href="https://aihuman.life">AIHuman.Life</a>还有其他AI工具可以体验吼~
            </p>
        </div>
    """
)


    fn=generate,
    run=generate,
    inputs=txt,
    outputs=out
    demo.launch(enable_queue=False, inline=True)