import torch
import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel, pipeline

tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model = GPT2LMHeadModel.from_pretrained('FredZhang7/anime-anything-promptgen-v2')

nlp = pipeline('text-generation', model=model, tokenizer=tokenizer)

def generate(prompt):
    output = ''
    outs = nlp(prompt, max_length=76, num_return_sequences=1, do_sample=True, repetition_penalty=1.2, temperature=0.7, top_k=3, early_stopping=True)
    for i in range(len(outs)):
        generated_text = str(outs[i]['generated_text']).replace('  ', '').rstrip(',')
        output += generated_text
        if i < len(outs)-1:
            output += '\n\n'
    return output

input_component = gr.Textbox(label="Prompt Idea", value="")
output_component = gr.Textbox(label="Extended Prompts")
examples = []
description = ""

gr.Interface(generate, inputs=input_component, outputs=output_component, examples=examples, title="Anime Prompt Gen", description=description).launch()