File size: 2,151 Bytes
98a1648 f81735b 98a1648 0067cc9 f87e5d7 98a1648 14e3b95 98a1648 ae3cb60 dfab003 ae3cb60 529eb10 98a1648 d82060c 98a1648 d82060c 98a1648 242e890 14e3b95 8dcc59b 242e890 98a1648 ae3cb60 242e890 d82060c 3d5ac88 98a1648 dfab003 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
def vipe_generate(text, model, tokenizer,device,do_sample,top_k=100, epsilon_cutoff=.00005, temperature=1):
#mark the text with special tokens
text=[tokenizer.eos_token + i + tokenizer.eos_token for i in text]
batch=tokenizer(text, padding=True, return_tensors="pt")
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
#how many new tokens to generate at max
max_prompt_length=50
generated_ids = model.generate(input_ids=input_ids,attention_mask=attention_mask, max_new_tokens=max_prompt_length, do_sample=do_sample,top_k=top_k, epsilon_cutoff=epsilon_cutoff, temperature=temperature)
#return only the generated prompts
pred_caps = tokenizer.batch_decode(generated_ids[:, -(generated_ids.shape[1] - input_ids.shape[1]):], skip_special_tokens=True)
return pred_caps[0]
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = GPT2LMHeadModel.from_pretrained('fittar/ViPE-M-CTX7')
model.to(device)
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
tokenizer.pad_token = tokenizer.eos_token
def generate(text):
result =vipe_generate([text],model,tokenizer,do_sample=True,device=device)
return result
examples = [
["brave, fantasy"],
["She felt like a flower in December"],
["2+2=4? hmm.."]
]
title = "ViPE: Visualize Pretty-much Everything"
description = 'ViPE is the first automated model for translating any arbitrary piece of text into a visualizable prompt. It helps any text-to-image model in figurative or non-lexical language visualizations. To learn more about the model, [click here](https://huggingface.co/fittar/ViPE-M-CTX7).<br>'
txt = gr.Textbox(lines=1, label="Arbitrary Input Text", placeholder="Initial text")
out = gr.Textbox(lines=4, label="Generated Prompt for Visualizations")
demo = gr.Interface(
fn =generate,
inputs=txt,
outputs=out,
examples=examples,
title=title,
description=description,
theme="default",
cache_examples="never"
)
demo.launch(enable_queue=True, debug=True) |