import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("crumb/bloom-560m-RLHF-SD2-prompter-aesthetic")
model = AutoModelForCausalLM.from_pretrained("crumb/bloom-560m-RLHF-SD2-prompter-aesthetic")
text2text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, num_workers=2)

def predict(text, max_length=64):
    text = text.strip()
    out_text = text2text_generator(text, max_new_tokens=max_length, 
                              eos_token_id = tokenizer.eos_token_id,
                              bos_token_id = tokenizer.bos_token_id,
                              pad_token_id = tokenizer.pad_token_id,
                              repetition_penalty = 1.05,
                         )[0]['generated_text']
    out_text = "<p>" + out_text + "</p>"
    out_text = out_text.replace(text, text + "<b><span>")
    out_text = out_text +  "</span></b>"
    out_text = out_text.replace("\n", "<br>")
    return out_text

iface = gr.Interface(
    fn=predict, 
    inputs=[
        gr.inputs.Textbox(lines=5, label="Input Text", default="Prompt: "),
        gr.inputs.Slider(minimum=1, maximum=64, default=32, step=1, label="Max Length"),
    ],
    outputs=gr.HTML(),
    description="use crumb/bloom-560m-RLHF-SD2-prompter-aesthetic to create prompts meant for [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base). tip: start with 'Prompt:'",
    examples=[[
            "Prompt: Pikachu,",
            32        
    ]
    ]
)

iface.launch()