import gradio as grad
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def load_prompter():
  prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
  tokenizer = AutoTokenizer.from_pretrained("gpt2")
  tokenizer.pad_token = tokenizer.eos_token
  tokenizer.padding_side = "left"
  return prompter_model, tokenizer

prompter_model, prompter_tokenizer = load_prompter()

def generate(plain_text):
    input_ids = prompter_tokenizer(plain_text.strip()+" Rephrase:", return_tensors="pt").input_ids
    eos_id = prompter_tokenizer.eos_token_id
    # Just use 1 beam and get 1 output, this is much, much, much faster than 8 beams and 8 outputs and we're only using the first.
    outputs = prompter_model.generate(input_ids, do_sample=False, max_new_tokens=75, eos_token_id=eos_id, pad_token_id=eos_id, length_penalty=-1.0)
    # Use [input_ids.shape[-1]:] because the decoded tokenised version of plain_text may have a different number of characters to the original
    res = prompter_tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
    return res

txt = grad.Textbox(lines=1, label="Initial Text", placeholder="Input Prompt")
out = grad.Textbox(lines=1, label="Optimized Prompt")
examples = ["A rabbit is wearing a space suit", "Several railroad tracks with one train passing by", "The roof is wet from the rain", "Cats dancing in a space club"]

grad.Interface(fn=generate,
               inputs=txt,
               outputs=out,
               title="Promptist Demo",
               description="Promptist is a prompt interface for Stable Diffusion v1-4 (https://huggingface.co./CompVis/stable-diffusion-v1-4) that optimizes user input into model-preferred prompts. The online demo at Hugging Face Spaces is using CPU, so slow generation speed would be expected. Please load the model locally with GPUs for faster generation.\n\nNote: This is a version with beam_size=1 while the original demo uses beam_size=8. So there would be a difference in terms of performance, but this demo is much faster. Many thanks to @HughPH for pointing out this improvement.",
               examples=examples,
               allow_flagging='never',
               cache_examples=False,
               theme="default").launch(enable_queue=True, debug=True)