MGM / app.py
wcy1122's picture
update demo
78fcda9
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, ack_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/MGM-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/MGM-13B-HD'
model_name = 'MGM-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':
print(generated_text)
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'
print(prompt, '---------')
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 'generate' in textbox.lower():
gen_image = 'Yes'
elif 'show me one idea of what i could make with this?' in textbox.lower() and imagebox is not None:
h, w = image.size
if h == 1505 and w == 1096:
gen_image = 'Yes'
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, gen_image) + (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='MGM') as demo:
gr.Markdown(title_markdown)
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], cache_examples=False)
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="MGM 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)
gr.Markdown(ack_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, gen_image] + 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, gen_image] + 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)