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='css') as demo:
 #   grad.HTML(
  #      """
   #         <div style="text-align: center; max-width: 650px; margin: 0 auto;">
     #         <div>
      #          <h5 style="font-weight: 900; font-size: 3rem; margin-bottom:20px;">
       #        Our  Prompt Generator!
        #        </h5><h5>Input a prompt idea and watch it preform a great prompt</h5>
         #     </div>
              
          #    </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="Enter a basic idea",
            ).style(
                container=False,
            )
            run = grad.Button("✨ Magic Prompt ✨").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>
     
        
    """
)


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