Spaces:
Running
Running
| 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') | |