fix

app.py

@@ -1,124 +1,76 @@
-import random
-import re
-
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import pipeline, set_seed
 
+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.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
+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
     )
-    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)
+    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,
-        max_new_tokens=75,
-        num_beams=8,
-        num_return_sequences=8,
-        length_penalty=-1.0
+        model_name='microsoft',
+        prompt_min_length=60,
+        prompt_max_length=75,
+        prompt_num_return_sequences=8,
+
 ):
-    result = generate_prompter(
-        bclip_text,
-        max_new_tokens,
-        num_beams,
-        num_return_sequences,
-        length_penalty
+    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
     )
-    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)
+    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)
 
 
 with gr.Blocks(title="Prompt生成器") 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输入)')
+
+            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('优化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')
-            output_w14 = gr.Textbox(label='W14的 Prompt')
+            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)")
@@ -126,38 +78,35 @@ with gr.Blocks(title="Prompt生成器") as block:
                 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')
+            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('SD优化')
-
-        with gr.Tab('文本生成'):
-            with gr.Row():
-                input_text = gr.Textbox(lines=6, label='你的想法', placeholder='在此输入内容...')
-                translate_output = gr.Textbox(lines=6, label='翻译结果(Prompt输入)')
+                img_prompter_btn = gr.Button('优化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('Prompt优化参数', open=True):
+                prompt_mode_name = gr.Radio(
+                    [
+                        'microsoft',
+                        'mj',
+                        'gpt2_650k',
+                    ],
+                    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, 6, 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', )
+                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(
                     [
@@ -185,22 +134,32 @@ with gr.Blocks(title="Prompt生成器") as block:
                 )
     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,
+        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]
     )
     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,
+        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(

utils/generator.py

@@ -0,0 +1,171 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, set_seed
+import random
+import re
+from .singleton import Singleton
+
+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 ('microsoft_model', 'microsoft_tokenizer'):
+            self.microsoft_model, self.microsoft_tokenizer = self.load_microsoft_model()
+
+        if item in ('mj_pipe',):
+            self.mj_pipe = self.load_mj_pipe()
+
+        if item in ('gpt2_650k_pipe',):
+            self.gpt2_650k_pipe = self.load_gpt2_650k_pipe()
+
+        return getattr(self, item)
+
+    @classmethod
+    def load_gpt2_650k_pipe(cls):
+
+        return pipeline('text-generation', model='Ar4ikov/gpt2-650k-stable-diffusion-prompt-generator')
+
+    @classmethod
+    def load_mj_pipe(cls):
+        return pipeline('text-generation', model='succinctly/text2image-prompt-generator')
+
+    @classmethod
+    def load_microsoft_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 rand_length(min_length: int = 60, max_length: int = 90) -> int:
+    if min_length > max_length:
+        return max_length
+
+    return random.randint(min_length, max_length)
+
+
+def generate_prompt(
+        plain_text,
+        min_length=60,
+        max_length=90,
+        num_return_sequences=8,
+        model_name='microsoft',
+):
+    if model_name == 'gpt2_650k':
+        return generate_prompt_gpt2_650k(
+            prompt=plain_text,
+            min_length=min_length,
+            max_length=max_length,
+            num_return_sequences=num_return_sequences,
+        )
+    elif model_name == 'mj':
+        return generate_prompt_mj(
+            text_in_english=plain_text,
+            num_return_sequences=num_return_sequences,
+            min_length=min_length,
+            max_length=max_length,
+        )
+    else:
+        return generate_prompt_microsoft(
+            plain_text=plain_text,
+            min_length=min_length,
+            max_length=max_length,
+            num_return_sequences=num_return_sequences,
+            num_beams=num_return_sequences,
+        )
+
+
+def generate_prompt_microsoft(
+        plain_text,
+        min_length=60,
+        max_length=90,
+        num_beams=8,
+        num_return_sequences=8,
+        length_penalty=-1.0
+) -> str:
+    input_ids = models.microsoft_tokenizer(plain_text.strip() + " Rephrase:", return_tensors="pt").input_ids
+    eos_id = models.microsoft_tokenizer.eos_token_id
+
+    outputs = models.microsoft_model.generate(
+        input_ids,
+        do_sample=False,
+        max_new_tokens=rand_length(min_length, max_length),
+        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.microsoft_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 generate_prompt_gpt2_650k(prompt: str, min_length=60, max_length: int = 255, num_return_sequences: int = 8) -> str:
+    def get_valid_prompt(text: str) -> str:
+        dot_split = text.split('.')[0]
+        n_split = text.split('\n')[0]
+
+        return {
+            len(dot_split) < len(n_split): dot_split,
+            len(n_split) > len(dot_split): n_split,
+            len(n_split) == len(dot_split): dot_split
+        }[True]
+
+    output = []
+    for _ in range(6):
+
+        output += [
+            get_valid_prompt(result['generated_text']) for result in
+            models.gpt2_650k_pipe(
+                prompt,
+                max_new_tokens=rand_length(min_length, max_length),
+                num_return_sequences=num_return_sequences
+            )
+        ]
+        output = list(set(output))
+        if len(output) >= num_return_sequences:
+            break
+
+    # valid_prompt = get_valid_prompt(models.gpt2_650k_pipe(prompt, max_length=max_length)[0]['generated_text'])
+    return "\n".join([o.strip() for o in output])
+
+
+def generate_prompt_mj(text_in_english: str, num_return_sequences: int = 8, min_length=60, max_length=90) -> str:
+    seed = random.randint(100, 1000000)
+    set_seed(seed)
+
+    result = ""
+    for _ in range(6):
+        sequences = models.mj_pipe(
+            text_in_english,
+            max_new_tokens=rand_length(min_length, max_length),
+            num_return_sequences=num_return_sequences
+        )
+        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
+    # return result, "\n".join(translate_en2zh(line) for line in result.split("\n") if len(line) > 0)