import shutil import subprocess import timm import spaces import io import base64 import torch import gradio as gr import os from PIL import Image import tempfile from huggingface_hub import snapshot_download from transformers import TextIteratorStreamer from threading import Thread from diffusers import StableDiffusionXLPipeline from minigemini.constants import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX from minigemini.mm_utils import process_images, load_image_from_base64, tokenizer_image_token from minigemini.conversation import default_conversation, conv_templates, SeparatorStyle, Conversation from minigemini.serve.gradio_web_server import function_markdown, tos_markdown, learn_more_markdown, title_markdown, block_css from minigemini.model.builder import load_pretrained_model # os.system('python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html') # os.system('pip install paddleocr>=2.0.1') # from paddleocr import PaddleOCR def download_model(repo_id): local_dir = os.path.join('./checkpoints', repo_id.split('/')[-1]) os.makedirs(local_dir) snapshot_download(repo_id=repo_id, local_dir=local_dir, local_dir_use_symlinks=False) if not os.path.exists('./checkpoints/'): os.makedirs('./checkpoints/') download_model('YanweiLi/Mini-Gemini-13B-HD') download_model('laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft-soup') device = "cuda" if torch.cuda.is_available() else "cpu" load_8bit = False load_4bit = False dtype = torch.float16 conv_mode = "vicuna_v1" model_path = './checkpoints/Mini-Gemini-13B-HD' model_name = 'Mini-Gemini-13B-HD' model_base = None tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name, load_8bit, load_4bit, device=device) diffusion_pipe = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16" ).to(device=device) if hasattr(model.config, 'image_size_aux'): if not hasattr(image_processor, 'image_size_raw'): image_processor.image_size_raw = image_processor.crop_size.copy() image_processor.crop_size['height'] = model.config.image_size_aux image_processor.crop_size['width'] = model.config.image_size_aux image_processor.size['shortest_edge'] = model.config.image_size_aux no_change_btn = gr.Button() enable_btn = gr.Button(interactive=True) disable_btn = gr.Button(interactive=False) def upvote_last_response(state): return ("",) + (disable_btn,) * 3 def downvote_last_response(state): return ("",) + (disable_btn,) * 3 def flag_last_response(state): return ("",) + (disable_btn,) * 3 def clear_history(): state = conv_templates[conv_mode].copy() return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5 def process_image(prompt, images): if images is not None and len(images) > 0: image_convert = images # Similar operation in model_worker.py image_tensor = process_images(image_convert, image_processor, model.config) image_grid = getattr(model.config, 'image_grid', 1) if hasattr(model.config, 'image_size_aux'): raw_shape = [image_processor.image_size_raw['height'] * image_grid, image_processor.image_size_raw['width'] * image_grid] image_tensor_aux = image_tensor image_tensor = torch.nn.functional.interpolate(image_tensor, size=raw_shape, mode='bilinear', align_corners=False) else: image_tensor_aux = [] if image_grid >= 2: raw_image = image_tensor.reshape(3, image_grid, image_processor.image_size_raw['height'], image_grid, image_processor.image_size_raw['width']) raw_image = raw_image.permute(1, 3, 0, 2, 4) raw_image = raw_image.reshape(-1, 3, image_processor.image_size_raw['height'], image_processor.image_size_raw['width']) if getattr(model.config, 'image_global', False): global_image = image_tensor if len(global_image.shape) == 3: global_image = global_image[None] global_image = torch.nn.functional.interpolate(global_image, size=[image_processor.image_size_raw['height'], image_processor.image_size_raw['width']], mode='bilinear', align_corners=False) # [image_crops, image_global] raw_image = torch.cat([raw_image, global_image], dim=0) image_tensor = raw_image.contiguous() image_tensor = image_tensor.unsqueeze(0) if type(image_tensor) is list: image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor] image_tensor_aux = [image.to(model.device, dtype=torch.float16) for image in image_tensor_aux] else: image_tensor = image_tensor.to(model.device, dtype=torch.float16) image_tensor_aux = image_tensor_aux.to(model.device, dtype=torch.float16) else: images = None image_tensor = None image_tensor_aux = [] image_tensor_aux = image_tensor_aux if len(image_tensor_aux) > 0 else None replace_token = DEFAULT_IMAGE_TOKEN if getattr(model.config, 'mm_use_im_start_end', False): replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) image_args = {"images": image_tensor, "images_aux": image_tensor_aux} return prompt, image_args @spaces.GPU def generate(state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens): prompt = state.get_prompt() images = state.get_images(return_pil=True) prompt, image_args = process_image(prompt, images) input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to("cuda:0") streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=30) max_new_tokens = 512 do_sample = True if temperature > 0.001 else False stop_str = state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2 thread = Thread(target=model.generate, kwargs=dict( inputs=input_ids, do_sample=do_sample, temperature=temperature, top_p=top_p, max_new_tokens=max_new_tokens, streamer=streamer, use_cache=True, **image_args )) thread.start() generated_text = '' for new_text in streamer: generated_text += new_text if generated_text.endswith(stop_str): generated_text = generated_text[:-len(stop_str)] state.messages[-1][-1] = generated_text yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) if gen_image == 'Yes' and '' in generated_text and '' in generated_text: common_neg_prompt = "out of frame, lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature" prompt = generated_text.split("")[1].split("")[0] generated_text = generated_text.split("")[0] + '\n' + 'Prompt: ' + prompt + '\n' torch.cuda.empty_cache() output_img = diffusion_pipe(prompt, negative_prompt=common_neg_prompt).images[0] buffered = io.BytesIO() output_img.save(buffered, format='JPEG') img_b64_str = base64.b64encode(buffered.getvalue()).decode() output = (generated_text, img_b64_str) state.messages[-1][-1] = output yield (state, state.to_gradio_chatbot(), "", None) + (enable_btn,) * 5 torch.cuda.empty_cache() @spaces.GPU def add_text(state, imagebox, textbox, image_process_mode, gen_image): if state is None: state = conv_templates[conv_mode].copy() if imagebox is not None: textbox = DEFAULT_IMAGE_TOKEN + '\n' + textbox image = Image.open(imagebox).convert('RGB') if gen_image == 'Yes': textbox = textbox + ' ' if imagebox is not None: textbox = (textbox, image, image_process_mode) state.append_message(state.roles[0], textbox) state.append_message(state.roles[1], None) yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) def delete_text(state, image_process_mode): state.messages[-1][-1] = None prev_human_msg = state.messages[-2] if type(prev_human_msg[1]) in (tuple, list): prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode) yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False) with gr.Blocks(title='Mini-Gemini') as demo: gr.Markdown(title_markdown) # state = default_conversation.copy() state = gr.State() with gr.Row(): with gr.Column(scale=3): imagebox = gr.Image(label="Input Image", type="filepath") image_process_mode = gr.Radio( ["Crop", "Resize", "Pad", "Default"], value="Default", label="Preprocess for non-square image", visible=False) gr.Examples(examples=[ ["./minigemini/serve/examples/monday.jpg", "Explain why this meme is funny, and generate a picture when the weekend coming."], ["./minigemini/serve/examples/woolen.png", "Show me one idea of what I could make with this?"], ["./minigemini/serve/examples/extreme_ironing.jpg", "What is unusual about this image?"], ["./minigemini/serve/examples/waterview.jpg", "What are the things I should be cautious about when I visit here?"], ], inputs=[imagebox, textbox]) with gr.Accordion("Function", open=True) as parameter_row: gen_image = gr.Radio(choices=['Yes', 'No'], value='No', interactive=True, label="Generate Image") with gr.Accordion("Parameters", open=False) as parameter_row: temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",) top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",) max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",) with gr.Column(scale=7): chatbot = gr.Chatbot( elem_id="chatbot", label="Mini-Gemini Chatbot", height=850, layout="panel", ) with gr.Row(): with gr.Column(scale=7): textbox.render() with gr.Column(scale=1, min_width=50): submit_btn = gr.Button(value="Send", variant="primary") with gr.Row(elem_id="buttons") as button_row: upvote_btn = gr.Button(value="👍 Upvote", interactive=False) downvote_btn = gr.Button(value="👎 Downvote", interactive=False) flag_btn = gr.Button(value="⚠️ Flag", interactive=False) regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False) clear_btn = gr.Button(value="🗑️ Clear", interactive=False) gr.Markdown(function_markdown) gr.Markdown(tos_markdown) gr.Markdown(learn_more_markdown) btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn] upvote_btn.click( upvote_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) downvote_btn.click( downvote_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) flag_btn.click( flag_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) clear_btn.click( clear_history, None, [state, chatbot, textbox, imagebox] + btn_list, queue=False ) regenerate_btn.click( delete_text, [state, image_process_mode], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) textbox.submit( add_text, [state, imagebox, textbox, image_process_mode, gen_image], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) submit_btn.click( add_text, [state, imagebox, textbox, image_process_mode, gen_image], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) demo.launch(debug=True)