init

app.py

@@ -0,0 +1,235 @@
+import random
+import re
+
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, set_seed
+
+from utils.image2text import git_image2text, w14_image2text, clip_image2text
+from utils.singleton import Singleton
+from utils.translate import en2zh as translate_en2zh
+from utils.translate import zh2en as translate_zh2en
+from utils.exif import get_image_info
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+@Singleton
+class Models(object):
+
+    def __getattr__(self, item):
+        if item in self.__dict__:
+            return getattr(self, item)
+
+        if item in ('big_model', 'big_processor'):
+            self.big_model, self.big_processor = self.load_image2text_model()
+
+        if item in ('prompter_model', 'prompter_tokenizer'):
+            self.prompter_model, self.prompter_tokenizer = self.load_prompter_model()
+
+        if item in ('text_pipe',):
+            self.text_pipe = self.load_text_generation_pipeline()
+
+        return getattr(self, item)
+
+    @classmethod
+    def load_text_generation_pipeline(cls):
+        return pipeline('text-generation', model='succinctly/text2image-prompt-generator')
+
+    @classmethod
+    def load_prompter_model(cls):
+        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
+
+
+models = Models.instance()
+
+
+def generate_prompter(plain_text, max_new_tokens=75, num_beams=8, num_return_sequences=8, length_penalty=-1.0):
+    input_ids = models.prompter_tokenizer(plain_text.strip() + " Rephrase:", return_tensors="pt").input_ids
+    eos_id = models.prompter_tokenizer.eos_token_id
+    outputs = models.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 = models.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 image_generate_prompter(
+        bclip_text,
+        w14_text,
+        max_new_tokens=75,
+        num_beams=8,
+        num_return_sequences=8,
+        length_penalty=-1.0
+):
+    result = generate_prompter(
+        bclip_text,
+        max_new_tokens,
+        num_beams,
+        num_return_sequences,
+        length_penalty
+    )
+    return "\n".join(["{},{}".format(line.strip(), w14_text.strip()) for line in result.split("\n") if len(line) > 0])
+
+
+def text_generate(text_in_english):
+    seed = random.randint(100, 1000000)
+    set_seed(seed)
+
+    result = ""
+    for _ in range(6):
+        sequences = models.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)
+
+
+with gr.Blocks(title="Prompt生成器") as block:
+    with gr.Column():
+
+        with gr.Tab('从图片中生成'):
+            with gr.Row():
+                input_image = gr.Image(type='pil')
+                exif_info = gr.HTML()
+            output_blip_or_clip = gr.Textbox(label='生成的 Prompt')
+            output_w14 = gr.Textbox(label='W14的 Prompt')
+
+            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)")
+            images_generate_prompter_output = gr.Textbox(lines=6, label='SD优化的 Prompt')
+            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('SD优化')
+
+        with gr.Tab('文本生成'):
+            with gr.Row():
+                input_text = gr.Textbox(lines=6, label='你的想法', placeholder='在此输入内容...')
+                translate_output = gr.Textbox(lines=6, label='翻译结果(Prompt输入)')
+
+            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.Accordion('SD优化参数', open=True):
+                max_new_tokens = gr.Slider(1, 512, 100, label='max_new_tokens', step=1)
+                nub_beams = gr.Slider(1, 30, 6, label='num_beams', step=1)
+                num_return_sequences = gr.Slider(1, 30, 6, label='num_return_sequences', step=1)
+                length_penalty = gr.Slider(-1.0, 1.0, -1.0, label='length_penalty')
+            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', )
+            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, max_new_tokens, nub_beams, num_return_sequences, length_penalty],
+        outputs=images_generate_prompter_output,
+    )
+    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_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')

requirements.txt

@@ -0,0 +1,6 @@
+transformers==4.27.4
+sentencepiece==0.1.97
+sacremoses==0.0.53
+clip-interrogator==0.6.0
+torch==2.0.0
+gradio==3.24.1