Promptist：基于强化学习的自动提示词优化

动态

• [演示版发布] 2022年12月：HuggingFace Space演示
• [模型发布] 2022年12月：模型加载说明
• [论文发布] 2022年12月：《文本到图像生成的提示词优化》论文链接

• 语言模型作为提示词接口，将用户输入优化为模型偏好的提示词

• 通过强化学习训练自动优化提示词的语言模型

加载Stable Diffusion v1.4预训练模型

您可以在https://huggingface.co/spaces/microsoft/Promptist尝试在线演示。

[注意] HuggingFace Space的在线演示使用CPU运行，生成速度较慢。建议在本地GPU环境加载模型以获得更快生成速度。

import gradio as grad
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def 加载提示器():
  提示器模型 = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
  分词器 = AutoTokenizer.from_pretrained("gpt2")
  分词器.pad_token = 分词器.eos_token
  分词器.padding_side = "left"
  return 提示器模型, 分词器

提示器模型, 提示器分词器 = 加载提示器()

def 生成优化提示(原始文本):
    输入序列 = 提示器分词器(原始文本.strip()+" 改写:", return_tensors="pt").input_ids
    结束符 = 提示器分词器.eos_token_id
    输出结果 = 提示器模型.generate(输入序列, do_sample=False, max_new_tokens=75, num_beams=8, num_return_sequences=8, eos_token_id=结束符, pad_token_id=结束符, length_penalty=-1.0)
    输出文本 = 提示器分词器.batch_decode(输出结果, skip_special_tokens=True)
    最终结果 = 输出文本[0].replace(原始文本+" 改写:", "").strip()
    return 最终结果

输入框 = grad.Textbox(lines=1, label="初始文本", placeholder="输入提示词")
输出框 = grad.Textbox(lines=1, label="优化后提示词")
示例 = ["穿太空服的兔子", "多条铁轨上行驶的火车", "被雨水打湿的屋顶", "太空俱乐部跳舞的猫咪"]

grad.Interface(fn=生成优化提示,
              inputs=输入框,
              outputs=输出框,
              title="Promptist演示",
              description="Promptist是Stable Diffusion v1-4(https://huggingface.co/CompVis/stable-diffusion-v1-4)的提示词优化接口，可将用户输入转化为模型偏好的提示词。",
              examples=示例,
              allow_flagging='never',
              cache_examples=False,
              theme="default").launch(enable_queue=True, debug=True)