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(theme='SebastianBravo/simci_css') as demo:
    grad.HTML(
        """
            <div style="text-align: center; max-width: 650px; margin: 0 auto;">
              <div>
                <h4>
                  Magic Prompt generator
                </h4>
             """
    )
    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="Enter a basic idea",
         
                container=False,
            )
          
            run = grad.Button("✨ Magic Prompt ✨")

         



        with grad.Row(variant="compact"):
            out = grad.Textbox(
                label="Generated Text",
                show_label=False,
                lines=5,
                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> Input a short prompt idea, get a uniq advanced prompt in seconds </p>
        </div>
    """
)


    fn=generate,
    run=generate,
    inputs=txt,
    outputs=out
    demo.launch()