From a7d98f494a002e009f410a855967a987ab2295e0 Mon Sep 17 00:00:00 2001 From: SillyLossy Date: Wed, 15 Feb 2023 01:38:21 +0200 Subject: [PATCH] Use BLIP to send a picture to model --- modules/bot_picture.py | 9 ++++++ modules/html_generator.py | 6 ++++ requirements.txt | 1 + server.py | 60 +++++++++++++++++++++++++++------------ 4 files changed, 58 insertions(+), 18 deletions(-) create mode 100644 modules/bot_picture.py diff --git a/modules/bot_picture.py b/modules/bot_picture.py new file mode 100644 index 00000000..bc957eb6 --- /dev/null +++ b/modules/bot_picture.py @@ -0,0 +1,9 @@ +from nataili_blip.model_manager import BlipModelManager +from nataili_blip.caption import Caption + +def load_model(): + model_name = "BLIP" + mm = BlipModelManager() + mm.download_model(model_name) + mm.load_blip(model_name) + return Caption(mm.loaded_models[model_name]["model"], mm.loaded_models[model_name]["device"]) \ No newline at end of file diff --git a/modules/html_generator.py b/modules/html_generator.py index ce7c6db7..ef1ba90e 100644 --- a/modules/html_generator.py +++ b/modules/html_generator.py @@ -217,6 +217,12 @@ def generate_chat_html(history, name1, name2, character): .body { } + + .body img { + max-width: 300px; + max-height: 300px; + border-radius: 20px; + } """ output = '' diff --git a/requirements.txt b/requirements.txt index 7420966f..1182f721 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ bitsandbytes==0.37.0 gradio==3.15.0 numpy safetensors==0.2.8 +nataili_blip git+https://github.com/huggingface/transformers diff --git a/server.py b/server.py index fd2bf7bf..47c7625f 100644 --- a/server.py +++ b/server.py @@ -23,6 +23,7 @@ from tqdm import tqdm from transformers import AutoConfig from transformers import AutoModelForCausalLM from transformers import AutoTokenizer +from io import BytesIO from modules.html_generator import * from modules.stopping_criteria import _SentinelTokenStoppingCriteria @@ -53,6 +54,7 @@ parser.add_argument('--listen', action='store_true', help='Make the web UI reach parser.add_argument('--listen-port', type=int, help='The listening port that the server will use.') parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.') parser.add_argument('--verbose', action='store_true', help='Print the prompts to the terminal.') +parser.add_argument('--picture', action='store_true', help='Adds an ability to send pictures in chat UI modes.') args = parser.parse_args() if (args.chat or args.cai_chat) and not args.no_stream: @@ -97,6 +99,10 @@ if args.deepspeed: ds_config = generate_ds_config(args.bf16, 1 * world_size, args.nvme_offload_dir) dschf = HfDeepSpeedConfig(ds_config) # Keep this object alive for the Transformers integration +if args.picture and (args.cai_chat or args.chat): + import modules.bot_picture as bot_picture + blip = bot_picture.load_model() + def load_model(model_name): print(f"Loading {model_name}...") t0 = time.time() @@ -561,8 +567,12 @@ def extract_message_from_reply(question, reply, current, other, check, extension return reply, next_character_found, substring_found -def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): +def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): original_text = text + + if args.picture and picture is not None: + text, original_text = generate_chat_picture(picture, name1, name2) + text = apply_extensions(text, "input") question = generate_chat_prompt(text, tokens, name1, name2, context, history_size) history['internal'].append(['', '']) @@ -573,12 +583,12 @@ def chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, history['internal'][-1] = [text, reply] history['visible'][-1] = [original_text, apply_extensions(reply, "output")] if not substring_found: - yield history['visible'] + yield history['visible'], None if next_character_found: break - yield history['visible'] + yield history['visible'], None -def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): +def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): question = generate_chat_prompt(text, tokens, name1, name2, context, history_size, impersonate=True) eos_token = '\n' if check else None for reply in generate_reply(question, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=eos_token, stopping_string=f"\n{name2}:"): @@ -589,20 +599,20 @@ def impersonate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, to break yield apply_extensions(reply, "output") -def cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): - for _history in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): - yield generate_chat_html(_history, name1, name2, character) +def cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): + for _history, _ in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): + yield generate_chat_html(_history, name1, name2, character), None -def regenerate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): +def regenerate_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): last = history['visible'].pop() history['internal'].pop() text = last[0] if args.cai_chat: - for i in cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): - yield i + for i, _ in cai_chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): + yield i, None else: - for i in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size): - yield i + for i, _ in chatbot_wrapper(text, tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size, picture): + yield i, None def remove_last_message(name1, name2): if not history['internal'][-1][0] == '<|BEGIN-VISIBLE-CHAT|>': @@ -791,6 +801,14 @@ def upload_your_profile_picture(img): img.save(Path(f'img_me.png')) print(f'Profile picture saved to "img_me.png"') +def generate_chat_picture(picture, name1, name2): + text = f'*{name1} sends {name2} a picture that contains the following: "{blip(picture)}"*' + buffer = BytesIO() + picture.save(buffer, format="JPEG") + img_str = base64.b64encode(buffer.getvalue()).decode('utf-8') + original_text = f'' + return text, original_text + # Global variables available_models = get_available_models() available_presets = get_available_presets() @@ -861,6 +879,9 @@ if args.chat or args.cai_chat: with gr.Row(): buttons["Send last reply to input"] = gr.Button("Send last reply to input") buttons["Replace last reply"] = gr.Button("Replace last reply") + if args.picture: + with gr.Row(): + picture_select = gr.Image(label="Send a picture", type='pil', display_label=True) with gr.Row(): with gr.Column(): @@ -906,14 +927,17 @@ if args.chat or args.cai_chat: if args.extensions is not None: create_extensions_block() - input_params = [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size_slider] + input_params = [textbox, max_new_tokens, do_sample, max_new_tokens, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, name1, name2, context, check, history_size_slider, picture_select] + output_params = [display, picture_select] if args.cai_chat: - gen_events.append(buttons["Generate"].click(cai_chatbot_wrapper, input_params, display, show_progress=args.no_stream, api_name="textgen")) - gen_events.append(textbox.submit(cai_chatbot_wrapper, input_params, display, show_progress=args.no_stream)) + gen_events.append(buttons["Generate"].click(cai_chatbot_wrapper, input_params, output_params, show_progress=args.no_stream, api_name="textgen")) + gen_events.append(textbox.submit(cai_chatbot_wrapper, input_params, output_params, show_progress=args.no_stream)) + picture_select.upload(cai_chatbot_wrapper, input_params, output_params, show_progress=args.no_stream) else: - gen_events.append(buttons["Generate"].click(chatbot_wrapper, input_params, display, show_progress=args.no_stream, api_name="textgen")) - gen_events.append(textbox.submit(chatbot_wrapper, input_params, display, show_progress=args.no_stream)) - gen_events.append(buttons["Regenerate"].click(regenerate_wrapper, input_params, display, show_progress=args.no_stream)) + gen_events.append(buttons["Generate"].click(chatbot_wrapper, input_params, output_params, show_progress=args.no_stream, api_name="textgen")) + gen_events.append(textbox.submit(chatbot_wrapper, input_params, output_params, show_progress=args.no_stream)) + picture_select.upload(chatbot_wrapper, input_params, output_params, show_progress=args.no_stream) + gen_events.append(buttons["Regenerate"].click(regenerate_wrapper, input_params, output_params, show_progress=args.no_stream)) gen_events.append(buttons["Impersonate"].click(impersonate_wrapper, input_params, textbox, show_progress=args.no_stream)) buttons["Send last reply to input"].click(send_last_reply_to_input, [], textbox, show_progress=args.no_stream)