init

app.py

@@ -0,0 +1,149 @@
+import random
+import re
+
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM
+from transformers import AutoProcessor
+from transformers import pipeline, set_seed
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+big_processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
+big_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")
+
+text_pipe = pipeline('text-generation', model='succinctly/text2image-prompt-generator')
+
+zh2en_model = AutoModelForSeq2SeqLM.from_pretrained('Helsinki-NLP/opus-mt-zh-en').eval()
+zh2en_tokenizer = AutoTokenizer.from_pretrained('Helsinki-NLP/opus-mt-zh-en')
+en2zh_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh").eval()
+en2zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
+
+
+def load_prompter():
+    prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+    return prompter_model, tokenizer
+
+
+prompter_model, prompter_tokenizer = load_prompter()
+
+
+def generate_prompter(plain_text, max_new_tokens=75, num_beams=8, num_return_sequences=8, length_penalty=-1.0):
+    input_ids = prompter_tokenizer(plain_text.strip() + " Rephrase:", return_tensors="pt").input_ids
+    eos_id = prompter_tokenizer.eos_token_id
+    outputs = prompter_model.generate(
+        input_ids,
+        do_sample=False,
+        max_new_tokens=max_new_tokens,
+        num_beams=num_beams,
+        num_return_sequences=num_return_sequences,
+        eos_token_id=eos_id,
+        pad_token_id=eos_id,
+        length_penalty=length_penalty
+    )
+    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    result = []
+    for output_text in output_texts:
+        result.append(output_text.replace(plain_text + " Rephrase:", "").strip())
+
+    return "\n".join(result)
+
+
+def translate_zh2en(text):
+    with torch.no_grad():
+        encoded = zh2en_tokenizer([text], return_tensors='pt')
+        sequences = zh2en_model.generate(**encoded)
+        return zh2en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
+
+
+def translate_en2zh(text):
+    with torch.no_grad():
+        encoded = en2zh_tokenizer([text], return_tensors="pt")
+        sequences = en2zh_model.generate(**encoded)
+        return en2zh_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
+
+
+def text_generate(text_in_english):
+    seed = random.randint(100, 1000000)
+    set_seed(seed)
+
+    result = ""
+    for _ in range(6):
+        sequences = text_pipe(text_in_english, max_length=random.randint(60, 90), num_return_sequences=8)
+        list = []
+        for sequence in sequences:
+            line = sequence['generated_text'].strip()
+            if line != text_in_english and len(line) > (len(text_in_english) + 4) and line.endswith(
+                    (':', '-', '—')) is False:
+                list.append(line)
+
+        result = "\n".join(list)
+        result = re.sub('[^ ]+\.[^ ]+', '', result)
+        result = result.replace('<', '').replace('>', '')
+        if result != '':
+            break
+    return result, "\n".join(translate_en2zh(line) for line in result.split("\n") if len(line) > 0)
+
+
+def get_prompt_from_image(input_image):
+    image = input_image.convert('RGB')
+    pixel_values = big_processor(images=image, return_tensors="pt").to(device).pixel_values
+
+    generated_ids = big_model.to(device).generate(pixel_values=pixel_values, max_length=50)
+    generated_caption = big_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    print(generated_caption)
+    return generated_caption
+
+
+with gr.Blocks() as block:
+    with gr.Column():
+        with gr.Tab('文本生成'):
+            with gr.Row():
+                input_text = gr.Textbox(lines=6, label='你的想法', placeholder='在此输入内容...')
+                translate_output = gr.Textbox(lines=6, label='翻译结果(Prompt输入)')
+
+            with gr.Accordion('SD优化参数设置', open=False):
+                max_new_tokens = gr.Slider(1, 255, 75, label='max_new_tokens', step=1)
+                nub_beams = gr.Slider(1, 30, 8, label='num_beams', step=1)
+                num_return_sequences = gr.Slider(1, 30, 8, label='num_return_sequences', step=1)
+                length_penalty = gr.Slider(-1.0, 1.0, -1.0, label='length_penalty')
+
+            generate_prompter_output = gr.Textbox(lines=6, label='SD优化的 Prompt')
+
+            output = gr.Textbox(lines=6, label='瞎编的 Prompt')
+            output_zh = gr.Textbox(lines=6, label='瞎编的 Prompt(zh)')
+            with gr.Row():
+                translate_btn = gr.Button('翻译')
+                generate_prompter_btn = gr.Button('SD优化')
+                gpt_btn = gr.Button('瞎编')
+
+        with gr.Tab('从图片中生成'):
+            with gr.Row():
+                input_image = gr.Image(type='pil')
+            img_btn = gr.Button('提交')
+            output_image = gr.Textbox(lines=6, label='生成的 Prompt')
+    translate_btn.click(
+        fn=translate_zh2en,
+        inputs=input_text,
+        outputs=translate_output
+    )
+    generate_prompter_btn.click(
+        fn=generate_prompter,
+        inputs=[translate_output, max_new_tokens, nub_beams, num_return_sequences, length_penalty],
+        outputs=generate_prompter_output
+    )
+    gpt_btn.click(
+        fn=text_generate,
+        inputs=translate_output,
+        outputs=[output, output_zh]
+    )
+    img_btn.click(
+        fn=get_prompt_from_image,
+        inputs=input_image,
+        outputs=output_image
+    )
+
+block.queue(max_size=64).launch(show_api=False, enable_queue=True, debug=True, share=False, server_name='0.0.0.0')

requirements.txt

@@ -0,0 +1,3 @@
+transformers==4.27.4
+#sentencepiece
+#sacremoses