antonovmaxim/text-generation-webui-space

1import base642import logging3import re4import time5from functools import partial6from io import BytesIO7 8import gradio as gr9import torch10 11from extensions.multimodal.multimodal_embedder import MultimodalEmbedder12from modules import shared13 14params = {15    "add_all_images_to_prompt": False,16    # device to run vision encoder on17    "vision_device": None,18    # bits to load vision encoder in, either 16 or 3219    "vision_bits": 32,20    # device to run multimodal projector on21    "projector_device": None,22    # multimodal projector bits, either 32 or 1623    "projector_bits": 3224}25 26 27# If 'state' is True, will hijack the next chat generation28input_hijack = {29    'state': False,30    'value': ["", ""]31}32 33 34# initialized in ui, so that params are loaded from settings35multimodal_embedder: MultimodalEmbedder = None36 37 38def add_chat_picture(picture, text, visible_text):39    # resize the image, so that shortest edge is at least 224 (size for CLIP), and at most 300 (to keep history manageable)40    max_hw, min_hw = max(picture.size), min(picture.size)41    aspect_ratio = max_hw / min_hw42    shortest_edge = int(max(300 / aspect_ratio, 224))43    longest_edge = int(shortest_edge * aspect_ratio)44    w = shortest_edge if picture.width < picture.height else longest_edge45    h = shortest_edge if picture.width >= picture.height else longest_edge46    picture = picture.resize((w, h))47 48    buffer = BytesIO()49    picture.save(buffer, format="JPEG")50    img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')51    image = f'<img src="data:image/jpeg;base64,{img_str}">'52 53    if '<image>' in text:54        text = text.replace('<image>', image)55    else:56        text = text + '\n' + image57 58    if visible_text == '' or visible_text is None:59        visible_text = text60    elif '<image>' in visible_text:61        visible_text = visible_text.replace('<image>', image)62    else:63        visible_text = visible_text + '\n' + image64 65    return text, visible_text66 67 68def custom_tokenized_length(prompt):69    return multimodal_embedder.len_in_tokens(prompt)70 71 72def tokenizer_modifier(state, prompt, input_ids, input_embeds):73    global params74    start_ts = time.time()75    image_match = re.search(r'<img src="data:image/jpeg;base64,[A-Za-z0-9+/=]+">', prompt)76 77    if image_match is None:78        return prompt, input_ids, input_embeds79 80    prompt, input_ids, input_embeds, total_embedded = multimodal_embedder.forward(prompt, state, params)81    logging.info(f'Embedded {total_embedded} image(s) in {time.time()-start_ts:.2f}s')82    return (prompt,83            input_ids.unsqueeze(0).to(shared.model.device, dtype=torch.int64),84            input_embeds.unsqueeze(0).to(shared.model.device, dtype=shared.model.dtype))85 86 87def ui():88    global multimodal_embedder89    multimodal_embedder = MultimodalEmbedder(params)90    with gr.Column():91        picture_select = gr.Image(label='Send a picture', type='pil')92        # The models don't seem to deal well with multiple images93        single_image_checkbox = gr.Checkbox(False, label='Embed all images, not only the last one')94    # Prepare the input hijack95    picture_select.upload(96        lambda picture: input_hijack.update({"state": True, "value": partial(add_chat_picture, picture)}),97        [picture_select],98        None99    )100    picture_select.clear(lambda: input_hijack.update({"state": False, "value": ["", ""]}), None, None)101    single_image_checkbox.change(lambda x: params.update({"add_all_images_to_prompt": x}), single_image_checkbox, None)102    shared.gradio['Generate'].click(lambda: None, None, picture_select)103    shared.gradio['textbox'].submit(lambda: None, None, picture_select)104