Views
No views yet
1from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
2
3device = "cuda" if torch.cuda.is_available() else "cpu"
4
5# Model checkpoint
6model_checkpoint = "Elixpo/promptPimp"
7
8# Tokenizer
9tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
10
11# Model
12model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
13
14enhancer = pipeline('text2text-generation',
15 model=model,
16 tokenizer=tokenizer,
17 repetition_penalty= 1.2,
18 device=device)
19
20max_target_length = 256
21prefix = "enhance prompt: "
22
23short_prompt = "beautiful house with text 'hello'"
24answer = enhancer(prefix + short_prompt, max_length=max_target_length)
25final_answer = answer[0]['generated_text']
26print(final_answer)
27
28# a two-story house with white trim, large windows on the second floor,
29# three chimneys on the roof, green trees and shrubs in front of the house,
30# stone pathway leading to the front door, text on the house reads "hello" in all caps,
31# blue sky above, shadows cast by the trees, sunlight creating contrast on the house's facade,
32# some plants visible near the bottom right corner, overall warm and serene atmosphere.