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

gpt2_pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', tokenizer='gpt2')
with open("ideas.txt", "r") as f:
    line = f.readlines()

def gerar_texto(texto_inicial):
    seed = random.randint(100, 1000000)
    set_seed(seed)

    # Se o texto inicial estiver vazio, selecione aleatoriamente uma linha do arquivo "ideas.txt" (em português)
    if texto_inicial == "":
        texto_inicial = line[random.randrange(0, len(line))].replace("\n", "").capitalize()
        texto_inicial = re.sub(r"[,:\-–.!;?_]", '', texto_inicial)

    response = gpt2_pipe(texto_inicial, max_length=(len(texto_inicial) + random.randint(60, 90)), num_return_sequences=4)
    response_list = []
    for x in response:
        resp = x['generated_text'].strip()
        if resp != texto_inicial and len(resp) > (len(texto_inicial) + 4) and not resp.endswith((":", "-", "—")):
            response_list.append(resp+'\n')

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

    if response_end != "":
        return response_end

txt = grad.Textbox(lines=1, label="Texto Inicial", placeholder="Digite o texto aqui")
out = grad.Textbox(lines=4, label="Prompts Gerados")

exemplos = []
for x in range(8):
    exemplos.append(line[random.randrange(0, len(line))].replace("\n", "").capitalize())

titulo = "Gerador de Prompt de Difusão Estável"
descricao = 'Este é um demo da série de modelos: "MagicPrompt", neste caso, focado em: "Difusão Estável". Para utilizá-lo, simplesmente insira seu texto ou clique em um dos exemplos. Para saber mais sobre o modelo, [clique aqui](https://huggingface.co/Gustavosta/MagicPrompt-Stable-Diffusion).<br>'

grad.Interface(fn=gerar_texto,
               inputs=txt,
               outputs=out,
               examples=exemplos,
               title=titulo,
               description=descricao,
               article='',
               allow_flagging='never',
               cache_examples=False,
               theme="default").launch(enable_queue=True, debug=True)