Spaces:
Runtime error
Runtime error
| 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 | |
| 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 '<h>' in generated_text and '</h>' 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("<h>")[1].split("</h>")[0] | |
| generated_text = generated_text.split("<h>")[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() | |
| 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 + ' <GEN>' | |
| 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) | |