import gradio as gr
import torch
import mdtex2html
from utils.exif import get_image_info
from utils.generator import generate_prompt
from utils.image2text import git_image2text, w14_image2text, clip_image2text
from utils.translate import en2zh as translate_en2zh
from utils.translate import zh2en as translate_zh2en
from utils.chatglm import chat2text
from utils.chatglm import models as chatglm_models
device = "cuda" if torch.cuda.is_available() else "cpu"
def text_generate_prompter(
plain_text,
model_name='microsoft',
prompt_min_length=60,
prompt_max_length=75,
prompt_num_return_sequences=8,
):
result = generate_prompt(
plain_text=plain_text,
model_name=model_name,
min_length=prompt_min_length,
max_length=prompt_max_length,
num_return_sequences=prompt_num_return_sequences
)
return result, "\n".join(translate_en2zh(line) for line in result.split("\n") if len(line) > 0)
def image_generate_prompter(
bclip_text,
w14_text,
model_name='microsoft',
prompt_min_length=60,
prompt_max_length=75,
prompt_num_return_sequences=8,
):
result = generate_prompt(
plain_text=bclip_text,
model_name=model_name,
min_length=prompt_min_length,
max_length=prompt_max_length,
num_return_sequences=prompt_num_return_sequences
)
prompter_list = ["{},{}".format(line.strip(), w14_text.strip()) for line in result.split("\n") if len(line) > 0]
prompter_zh_list = [
"{},{}".format(translate_en2zh(line.strip()), translate_en2zh(w14_text.strip())) for line in
result.split("\n") if len(line) > 0
]
return "\n".join(prompter_list), "\n".join(prompter_zh_list)
def translate_input(text: str, chatglm_text: str) -> str:
if chatglm_text is not None and len(chatglm_text) > 0:
return translate_zh2en(chatglm_text)
return translate_zh2en(text)
with gr.Blocks(title="Prompt生成器") as block:
with gr.Column():
with gr.Tab('Chat'):
def revise(history, latest_message):
history[-1] = (history[-1][0], latest_message)
return history, ''
def revoke(history):
if len(history) >= 1:
history.pop()
return history
def interrupt(allow_generate):
allow_generate[0] = False
def reset_state():
return [], []
with gr.Row():
with gr.Column(scale=4):
chatbot = gr.Chatbot(elem_id="chat-box", show_label=False).style(height=800)
with gr.Column(scale=1):
with gr.Row():
max_length = gr.Slider(32, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
top_p = gr.Slider(0.01, 1, value=0.7, step=0.01, label="Top P", interactive=True)
temperature = gr.Slider(0.01, 5, value=0.95, step=0.01, label="Temperature", interactive=True)
with gr.Row():
query = gr.Textbox(show_label=False, placeholder="Prompts", lines=4).style(container=False)
generate_button = gr.Button("生成")
with gr.Row():
continue_message = gr.Textbox(
show_label=False, placeholder="Continue message", lines=2).style(container=False)
continue_btn = gr.Button("续写")
revise_message = gr.Textbox(
show_label=False, placeholder="Revise message", lines=2).style(container=False)
revise_btn = gr.Button("修订")
revoke_btn = gr.Button("撤回")
interrupt_btn = gr.Button("终止生成")
reset_btn = gr.Button("清空")
history = gr.State([])
allow_generate = gr.State([True])
blank_input = gr.State("")
reset_btn.click(reset_state, outputs=[chatbot, history], show_progress=True)
generate_button.click(
chatglm_models.chatglm.predict_continue,
inputs=[query, blank_input, max_length, top_p, temperature, allow_generate, history],
outputs=[chatbot, query]
)
revise_btn.click(revise, inputs=[history, revise_message], outputs=[chatbot, revise_message])
revoke_btn.click(revoke, inputs=[history], outputs=[chatbot])
continue_btn.click(
chatglm_models.chatglm.predict_continue,
inputs=[query, continue_message, max_length, top_p, temperature, allow_generate, history],
outputs=[chatbot, query, continue_message]
)
interrupt_btn.click(interrupt, inputs=[allow_generate])
with gr.Tab('文本生成'):
with gr.Row():
input_text = gr.Textbox(lines=6, label='你的想法', placeholder='在此输入内容...')
chatglm_output = gr.Textbox(lines=6, label='ChatGLM', placeholder='在此输入内容...')
translate_output = gr.Textbox(lines=6, label='翻译结果(Prompt输入)')
output = gr.Textbox(lines=6, label='优化的 Prompt')
output_zh = gr.Textbox(lines=6, label='优化的 Prompt(zh)')
with gr.Row():
chatglm_btn = gr.Button('召唤ChatGLM')
translate_btn = gr.Button('翻译')
generate_prompter_btn = gr.Button('优化Prompt')
with gr.Tab('从图片中生成'):
with gr.Row():
input_image = gr.Image(type='pil')
exif_info = gr.HTML()
output_blip_or_clip = gr.Textbox(label='生成的 Prompt', lines=4)
output_w14 = gr.Textbox(label='W14的 Prompt', lines=4)
with gr.Accordion('W14', open=False):
w14_raw_output = gr.Textbox(label="Output (raw string)")
w14_booru_output = gr.Textbox(label="Output (booru string)")
w14_rating_output = gr.Label(label="Rating")
w14_characters_output = gr.Label(label="Output (characters)")
w14_tags_output = gr.Label(label="Output (tags)")
output_img_prompter = gr.Textbox(lines=6, label='优化的 Prompt')
output_img_prompter_zh = gr.Textbox(lines=6, label='优化的 Prompt(zh)')
with gr.Row():
img_exif_btn = gr.Button('EXIF')
img_blip_btn = gr.Button('BLIP图片转描述')
img_w14_btn = gr.Button('W14图片转描述')
img_clip_btn = gr.Button('CLIP图片转描述')
img_prompter_btn = gr.Button('优化Prompt')
with gr.Tab('参数设置'):
with gr.Accordion('Prompt优化参数', open=True):
prompt_mode_name = gr.Radio(
[
'microsoft',
'mj',
'gpt2_650k',
'gpt_neo_125m',
],
value='gpt2_650k',
label='model_name'
)
prompt_min_length = gr.Slider(1, 512, 100, label='min_length', step=1)
prompt_max_length = gr.Slider(1, 512, 200, label='max_length', step=1)
prompt_num_return_sequences = gr.Slider(1, 30, 8, label='num_return_sequences', step=1)
with gr.Accordion('BLIP参数', open=True):
blip_max_length = gr.Slider(1, 512, 100, label='max_length', step=1)
with gr.Accordion('CLIP参数', open=True):
clip_mode_type = gr.Radio(['best', 'classic', 'fast', 'negative'], value='best', label='mode_type')
clip_model_name = gr.Radio(['vit_h_14', 'vit_l_14', ], value='vit_h_14', label='model_name')
with gr.Accordion('WD14参数', open=True):
image2text_model = gr.Radio(
[
"SwinV2",
"ConvNext",
"ConvNextV2",
"ViT",
],
value="ConvNextV2",
label="Model"
)
general_threshold = gr.Slider(
0,
1,
step=0.05,
value=0.35,
label="General Tags Threshold",
)
character_threshold = gr.Slider(
0,
1,
step=0.05,
value=0.85,
label="Character Tags Threshold",
)
img_prompter_btn.click(
fn=image_generate_prompter,
inputs=[
output_blip_or_clip,
output_w14,
prompt_mode_name,
prompt_min_length,
prompt_max_length,
prompt_num_return_sequences,
],
outputs=[output_img_prompter, output_img_prompter_zh]
)
chatglm_btn.click(
fn=chatglm_models.chatglm.generator_image_text,
inputs=input_text,
outputs=chatglm_output,
)
translate_btn.click(
fn=translate_input,
inputs=[input_text, chatglm_output],
outputs=translate_output
)
generate_prompter_btn.click(
fn=text_generate_prompter,
inputs=[
translate_output,
prompt_mode_name,
prompt_min_length,
prompt_max_length,
prompt_num_return_sequences,
],
outputs=[output, output_zh]
)
img_w14_btn.click(
fn=w14_image2text,
inputs=[input_image, image2text_model, general_threshold, character_threshold],
outputs=[
output_w14,
w14_raw_output,
w14_booru_output,
w14_rating_output,
w14_characters_output,
w14_tags_output
]
)
img_blip_btn.click(
fn=git_image2text,
inputs=[input_image, blip_max_length],
outputs=output_blip_or_clip
)
img_clip_btn.click(
fn=clip_image2text,
inputs=[input_image, clip_mode_type, clip_model_name],
outputs=output_blip_or_clip
)
img_exif_btn.click(
fn=get_image_info,
inputs=input_image,
outputs=exif_info
)
block.queue(max_size=64).launch(show_api=False, enable_queue=True, debug=True, share=False, server_name='0.0.0.0')