基础模型: google-t5/t5-base 数据集:

• gokaygokay/prompt-enhancer-dataset 语言:
• en 库名称: transformers 许可证: apache-2.0 管道标签: text2text-generation

from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

device = "cuda" if torch.cuda.is_available() else "cpu"

# 模型检查点
model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"

# 分词器
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

# 模型
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

enhancer = pipeline('text2text-generation',
                    model=model,
                    tokenizer=tokenizer,
                    repetition_penalty= 1.2,
                    device=device)

max_target_length = 256
prefix = "enhance prompt: "

short_prompt = "beautiful house with text 'hello'"
answer = enhancer(prefix + short_prompt, max_length=max_target_length)
final_answer = answer[0]['generated_text']
print(final_answer)

# 一栋两层楼的房子，带有白色装饰，二楼有大窗户，
# 屋顶上有三个烟囱，房子前面有绿色的树木和灌木，
# 石头小径通向大门，房子上的文字是大写的“hello”，
# 上方是蓝天，树木投下阴影，阳光在房子的外立面上形成对比，
# 右下角附近可以看到一些植物，整体氛围温暖而宁静。