Update app.py

app.py

@@ -2,6 +2,8 @@ from transformers import pipeline
 import gradio as gr
 import random
 import paddlehub as hub
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from loguru import logger
 
 language_translation_model = hub.Module(directory=f'./baidu_translate')
@@ -23,10 +25,26 @@ def getTextTrans(text, source='zh', target='en'):
         
 extend_prompt_pipe = pipeline('text-generation', model='yizhangliu/prompt-extend', max_length=77, pad_token_id=0)
 
+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 extend_prompt_microsoft(in_text):
+    input_ids = prompter_tokenizer(in_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=75, num_beams=8, num_return_sequences=8, eos_token_id=eos_id, pad_token_id=eos_id, length_penalty=-1.0)
+    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    res = output_texts[0].replace(in_text+" Rephrase:", "").strip()
+    return res
+    
 space_ids = {
-            "spaces/stabilityai/stable-diffusion":"SD 2.1",
-            "spaces/runwayml/stable-diffusion-v1-5":"SD 1.5",
-            "spaces/stabilityai/stable-diffusion-1":"SD 1.0",
+            "spaces/stabilityai/stable-diffusion": "SD 2.1",
+            "spaces/runwayml/stable-diffusion-v1-5": "SD 1.5",
+            "spaces/stabilityai/stable-diffusion-1": "SD 1.0",
+            "spaces/IDEA-CCNL/Taiyi-Stable-Diffusion-Chinese": "Taiyi(太乙)",
             }
 
 tab_actions = []
@@ -34,6 +52,7 @@ tab_titles = []
 
 thanks_info = "Thanks: "
 thanks_info += "[<a style='display:inline-block' href='https://huggingface.co/spaces/daspartho/prompt-extend' _blank><font style='color:blue;weight:bold;'>prompt-extend</font></a>]"
+thanks_info += "[<a style='display:inline-block' href='https://huggingface.co/spaces/microsoft/Promptist' _blank><font style='color:blue;weight:bold;'>Promptist</font></a>]"
 
 for space_id in space_ids.keys():
     print(space_id, space_ids[space_id])
@@ -82,22 +101,20 @@ start_work = """async() => {
     
     if (typeof window['gradioEl'] === 'undefined') {
         window['gradioEl'] = gradioEl;
-        
         tabitems = window['gradioEl'].querySelectorAll('.tabitem');
+        tabitems_title = window['gradioEl'].querySelectorAll('#tab_demo')[0].children[0].children[0].children;
         
-        for (var i = 0; i < tabitems.length; i++) {    
-            if ([0, 1, 2].includes(i)) {
+        for (var i = 0; i < tabitems.length; i++) {   
+            if (tabitems_title[i].innerText.indexOf('SD') >= 0) {
                 tabitems[i].childNodes[0].children[0].style.display='none';
                 for (var j = 0; j < tabitems[i].childNodes[0].children[1].children.length; j++) {
                     if (j != 1) {
                         tabitems[i].childNodes[0].children[1].children[j].style.display='none';
                     }
                 }
-            } else {
-                tabitems[i].childNodes[0].children[0].style.display='none';
-                tabitems[i].childNodes[0].children[1].style.display='none';
-                tabitems[i].childNodes[0].children[2].children[0].style.display='none';
-                tabitems[i].childNodes[0].children[3].style.display='none';                
+            } else if (tabitems_title[i].innerText.indexOf('Taiyi') >= 0) {
+                tabitems[3].children[0].children[0].children[1].style.display='none';
+                tabitems[i].children[0].children[0].children[0].children[0].children[1].style.display='none';
             }            
         }  
         
@@ -106,8 +123,12 @@ start_work = """async() => {
         tab_demo.setAttribute('style', 'height: 100%;');
         const page1 = window['gradioEl'].querySelectorAll('#page_1')[0];
         const page2 = window['gradioEl'].querySelectorAll('#page_2')[0];
-        window['gradioEl'].querySelector('#input_col1_row2').children[0].setAttribute('style', 'min-width:0px;width:50%;');
-        window['gradioEl'].querySelector('#input_col1_row2').children[1].setAttribute('style', 'min-width:0px;width:50%;');
+        
+        btns_1 = window['gradioEl'].querySelector('#input_col1_row2').children;
+        btns_1_split = 100 / btns_1.length;
+        for (var i = 0; i < btns_1.length; i++) {
+            btns_1[i].setAttribute('style', 'min-width:0px;width:' + btns_1_split + '%;');
+        }
         page1.style.display = "none";
         page2.style.display = "block";    
         window['prevPrompt'] = '';
@@ -121,20 +142,25 @@ start_work = """async() => {
                             window['doCheckPrompt'] = 1;
                             window['prevPrompt'] = text_value;
                             tabitems = window['gradioEl'].querySelectorAll('.tabitem');
-                            for (var i = 0; i < tabitems.length; i++) {   
-                                if ([0, 1, 2].includes(i)) {
+                            for (var i = 0; i < tabitems.length; i++) {  
+                                inputText = null;
+                                if (tabitems_title[i].innerText.indexOf('SD') >= 0) {
+                                    text_value = window['gradioEl'].querySelectorAll('#prompt_work')[0].querySelectorAll('textarea')[0].value;
                                     inputText = tabitems[i].children[0].children[1].children[0].querySelectorAll('.gr-text-input')[0];
-                                } else {
-                                    inputText = tabitems[i].childNodes[0].children[2].children[0].children[0].querySelectorAll('.gr-text-input')[0];
+                                } else if (tabitems_title[i].innerText.indexOf('Taiyi') >= 0) {
+                                    text_value = window['gradioEl'].querySelectorAll('#prompt_work_zh')[0].querySelectorAll('textarea')[0].value;
+                                    inputText = tabitems[i].children[0].children[0].children[1].querySelectorAll('.gr-text-input')[0];
+                                }
+                                if (inputText) {
+                                    setNativeValue(inputText, text_value);
+                                    inputText.dispatchEvent(new Event('input', { bubbles: true }));
                                 }
-                                setNativeValue(inputText, text_value);
-                                inputText.dispatchEvent(new Event('input', { bubbles: true }));
                             }
                            
                             setTimeout(function() {
                                 btns = window['gradioEl'].querySelectorAll('button');
                                 for (var i = 0; i < btns.length; i++) {
-                                    if (['Generate image','Run'].includes(btns[i].innerText)) {
+                                    if (['Generate image','Run', '生成图像(Generate)'].includes(btns[i].innerText)) {
                                         btns[i].click();                
                                     }
                                 }
@@ -150,25 +176,40 @@ start_work = """async() => {
     return false;
 }"""
 
-def prompt_extend(prompt):    
+def prompt_extend(prompt, PM):    
     prompt_en = getTextTrans(prompt, source='zh', target='en')
-    extend_prompt_en = extend_prompt_pipe(prompt_en+',', num_return_sequences=1)[0]["generated_text"]    
+    if PM == 1:
+        extend_prompt_en = extend_prompt_pipe(prompt_en+',', num_return_sequences=1)[0]["generated_text"]   
+    else:
+        extend_prompt_en = extend_prompt_microsoft(prompt_en)
+        
     if (prompt != prompt_en):
-        logger.info(f"extend_prompt__1__")
+        logger.info(f"extend_prompt__1_[{PM}]_")
         extend_prompt_out = getTextTrans(extend_prompt_en, source='en', target='zh')
     else:
-        logger.info(f"extend_prompt__2__")
+        logger.info(f"extend_prompt__2_[{PM}]_")
         extend_prompt_out = extend_prompt_en
 
     return extend_prompt_out
 
+def prompt_extend_1(prompt):  
+    extend_prompt_out = prompt_extend(prompt, 1)
+    return extend_prompt_out
+
+def prompt_extend_2(prompt):  
+    extend_prompt_out = prompt_extend(prompt, 2)
+    return extend_prompt_out
+
 def prompt_draw(prompt):
     prompt_en = getTextTrans(prompt, source='zh', target='en')
     if (prompt != prompt_en):
         logger.info(f"draw_prompt______1__")
+        prompt_zh = prompt
     else:
         logger.info(f"draw_prompt______2__")
-    return prompt_en
+        prompt_zh = getTextTrans(prompt, source='en', target='zh')
+        
+    return prompt_en, prompt_zh
         
 with gr.Blocks(title='Text-to-Image') as demo:
     with gr.Group(elem_id="page_1", visible=True) as page_1:
@@ -183,21 +224,25 @@ with gr.Blocks(title='Text-to-Image') as demo:
                     with gr.Row(elem_id="input_col1_row1"):
                         prompt_input0 = gr.Textbox(lines=2, label="Original prompt", visible=True)
                     with gr.Row(elem_id="input_col1_row2"):
-                        with gr.Column(elem_id="input_col1_row2_col1"):
+                        with gr.Column(elem_id="input_col1_row2_col0"):
                             draw_btn_0 = gr.Button(value = "Generate(original)", elem_id="draw-btn-0")
+                        with gr.Column(elem_id="input_col1_row2_col1"):
+                            extend_btn_1 = gr.Button(value = "Extend_1",elem_id="extend-btn-1")   
                         with gr.Column(elem_id="input_col1_row2_col2"):
-                            extend_btn = gr.Button(value = "Extend prompt",elem_id="extend-btn")                    
+                            extend_btn_2 = gr.Button(value = "Extend_2",elem_id="extend-btn-2")                              
                 with gr.Column(id="input_col2"):
                     prompt_input1 = gr.Textbox(lines=2, label="Extend prompt", visible=True)
                     draw_btn_1 = gr.Button(value = "Generate(extend)", elem_id="draw-btn-1")
                     prompt_work = gr.Textbox(lines=1, label="prompt_work", elem_id="prompt_work", visible=False)
+                    prompt_work_zh = gr.Textbox(lines=1, label="prompt_work_zh", elem_id="prompt_work_zh", visible=False)
             with gr.Row(elem_id='tab_demo', visible=True).style(height=200):
                 tab_demo = gr.TabbedInterface(tab_actions, tab_titles) 
             with gr.Row():
-                gr.HTML(f"<p>{thanks_info}</p>")
+                gr.HTML(f"<p>{thanks_info}</p>") 
 
-            extend_btn.click(fn=prompt_extend, inputs=[prompt_input0], outputs=[prompt_input1])
-            draw_btn_0.click(fn=prompt_draw, inputs=[prompt_input0], outputs=[prompt_work])
-            draw_btn_1.click(fn=prompt_draw, inputs=[prompt_input1], outputs=[prompt_work])
+            extend_btn_1.click(fn=prompt_extend_1, inputs=[prompt_input0], outputs=[prompt_input1])
+            extend_btn_2.click(fn=prompt_extend_2, inputs=[prompt_input0], outputs=[prompt_input1])
+            draw_btn_0.click(fn=prompt_draw, inputs=[prompt_input0], outputs=[prompt_work, prompt_work_zh])
+            draw_btn_1.click(fn=prompt_draw, inputs=[prompt_input1], outputs=[prompt_work, prompt_work_zh])
         
 demo.launch()