Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import re | |
| from PIL import Image | |
| from tqdm import tqdm | |
| from train.scripts.generate_lm_multiple import gen_sequence, build_visorgpt | |
| from utils.seq2coord import gen_cond_mask | |
| from visor_gligen.gligen_inference_box import gligen_infer, build_gligen_model | |
| from visor_controlnet.gradio_pose2image_v2 import control_infer, build_control_model, build_controlv11_model | |
| # init models | |
| visorgpt_config_path = 'train/models/gpt2/config.json' | |
| visorgpt_model_path = 'demo/ckpts/visorgpt/visorgpt_dagger_ta_tb.pt' | |
| visorgpt_vocab_path = 'train/models/google_uncased_en_coord_vocab.txt' | |
| # control_model_path = 'demo/ckpts/controlnet/control_sd15_openpose.pth' | |
| control_model_path = 'demo/ckpts/controlnet/control_v11p_sd15_openpose.pth' # v1.1 | |
| control_sd_path = 'demo/ckpts/controlnet/v1-5-pruned-emaonly.safetensors' | |
| control_model_config = 'demo/ckpts/controlnet/cldm_v15.yaml' | |
| gligen_model_path = 'demo/ckpts/gligen/diffusion_pytorch_model_box.bin' | |
| visorgpt_args, visorgpt_model = build_visorgpt(model_config=visorgpt_config_path, | |
| model_path=visorgpt_model_path, | |
| vocab_path=visorgpt_vocab_path) | |
| control_model, ddim_sampler = build_controlv11_model(model_path=control_model_path, | |
| sd_path=control_sd_path, | |
| config_path=control_model_config) | |
| # build gligen model | |
| g_model, g_autoencoder, g_text_encoder, g_diffusion, \ | |
| g_config, g_grounding_tokenizer_input = build_gligen_model(ckpt=gligen_model_path) | |
| # maximum number of instances | |
| max_num_keypoint = 16 | |
| max_num_bbox = 16 | |
| max_num_mask = 8 | |
| def generate_sequence(gen_type, | |
| data_type, | |
| instance_size, | |
| num_instance, | |
| object_name_inbox): | |
| ctn = True | |
| if gen_type == 'key point': | |
| num_keypoint = 18 | |
| if num_instance > max_num_keypoint: | |
| num_instance = max_num_keypoint | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, str(num_instance), str(num_keypoint)]) + ' ; [person' | |
| elif gen_type == 'box' or gen_type == 'mask': | |
| if not object_name_inbox.strip(): | |
| if gen_type == 'mask': | |
| object_name_inbox = "bottle; cup" | |
| else: | |
| if data_type == 'object centric': | |
| object_name_inbox = "great white shark" | |
| else: | |
| object_name_inbox = "person; frisbee" | |
| num_keypoint = 0 | |
| if gen_type == 'mask': | |
| if num_instance > max_num_mask: | |
| num_instance = max_num_mask | |
| if gen_type == 'box': | |
| if num_instance > max_num_bbox: | |
| num_instance = max_num_bbox | |
| if data_type == 'object centric': | |
| num_instance = 1 | |
| objects = ', '.join(object_name_inbox.strip().split(";")) | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, | |
| str(num_instance), str(num_keypoint)]) + '; ' + objects | |
| if len(object_name_inbox.split(';')) > num_instance: | |
| return { | |
| raw_sequence: gr.update( | |
| value="The umber of category names should be less than the number of instances, please try again :)", | |
| visible=True) | |
| } | |
| print("input prompt: \n", seq_prompt) | |
| sequence = gen_sequence(visorgpt_args, visorgpt_model, seq_prompt) | |
| assert isinstance(sequence, list) | |
| try: | |
| cond_mask, cond_json = gen_cond_mask(sequence, ctn) | |
| if gen_type == 'key point': | |
| ori_sequence = cond_json[2]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'box': | |
| ori_sequence = cond_json[0]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'mask': | |
| ori_sequence = cond_json[1]['sequences'][0][0] + '[SEP]' | |
| except: | |
| cond_mask, cond_json = gen_cond_mask(sequence, not ctn) | |
| if gen_type == 'key point': | |
| ori_sequence = cond_json[2]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'box': | |
| ori_sequence = cond_json[0]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'mask': | |
| ori_sequence = cond_json[1]['sequences'][0][0] + '[SEP]' | |
| ret_img = Image.fromarray(cond_mask) | |
| if not gen_type == 'mask': | |
| return { | |
| result_gallery: [ret_img], | |
| raw_sequence: gr.update(value=ori_sequence, visible=True), | |
| images_button: gr.update(visible=True), | |
| text_container: cond_json, | |
| sequence_container: ori_sequence | |
| } | |
| else: | |
| return { | |
| result_gallery: [ret_img], | |
| raw_sequence: gr.update(value=ori_sequence, visible=True), | |
| images_button: gr.update(visible=False), | |
| text_container: cond_json, | |
| sequence_container: ori_sequence | |
| } | |
| def add_contents(gen_type, | |
| data_type, | |
| instance_size, | |
| num_instance, | |
| object_name_inbox, | |
| num_continuous_gen, | |
| global_seq): | |
| ctn = True | |
| if gen_type == 'key point': | |
| num_keypoint = 18 | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, str(num_instance), str(num_keypoint)]) + ' ; [person' | |
| if num_continuous_gen: | |
| ctn = True | |
| cur_instance = int(global_seq.split(';')[3].strip()) | |
| new_number = cur_instance + num_continuous_gen | |
| if new_number > max_num_keypoint: | |
| new_number = max_num_keypoint | |
| # prompt type a | |
| if global_seq.split(';')[5].find('[') == -1: | |
| global_seq = global_seq.replace('[CLS]', '').replace('[SEP]', '') | |
| objects = re.findall(re.compile(r'[\[](.*?)[]]', re.S), global_seq) | |
| objects = ' '.join(['[ person' + x + ']' for x in objects]) | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, str(new_number), str(num_keypoint), objects]) | |
| # prompt type b | |
| else: | |
| global_seq = global_seq.replace('[CLS]', '').replace('[SEP]', '') | |
| seq_list = global_seq.split(';') | |
| seq_list[3] = str(new_number) | |
| seq_prompt = ';'.join(seq_list) | |
| elif gen_type == 'box' or gen_type == 'mask': | |
| num_keypoint = 0 | |
| if data_type == 'object centric': | |
| num_instance = 1 | |
| objects = ', '.join(object_name_inbox.strip().split(";")) | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, | |
| str(num_instance), str(num_keypoint)]) + '; ' + objects | |
| if len(object_name_inbox.split(';')) > num_instance: | |
| return { | |
| raw_sequence: gr.update(value=f"The umber of category names should be less than the number of instances, please try again :)", visible=True) | |
| } | |
| if num_continuous_gen: | |
| cur_instance = int(global_seq.split(';')[3].strip()) | |
| new_number = cur_instance + num_continuous_gen | |
| if gen_type == 'mask': | |
| if new_number > max_num_mask: | |
| new_number = max_num_mask | |
| if gen_type == 'box': | |
| if new_number > max_num_bbox: | |
| new_number = max_num_bbox | |
| # prompt type a | |
| if global_seq.split(';')[5].find('[') == -1: | |
| global_seq = global_seq.replace('[CLS]', '').replace('[SEP]', '') | |
| coords = re.findall(re.compile(r'[\[](.*?)[]]', re.S), global_seq) | |
| objects = global_seq.split(';')[5].split(',') | |
| objects = ' '.join(['[ ' + objects[i] + coords[i] + ']' for i in range(len(coords))]) | |
| seq_prompt = '; '.join([gen_type, data_type, instance_size, str(new_number), str(num_keypoint), objects]) | |
| # prompt type b | |
| else: | |
| global_seq = global_seq.replace('[CLS]', '').replace('[SEP]', '') | |
| seq_list = global_seq.split(';') | |
| seq_list[3] = str(new_number) | |
| seq_prompt = ';'.join(seq_list) | |
| # import ipdb;ipdb.set_trace() | |
| print("input prompt: \n", seq_prompt) | |
| with torch.no_grad(): | |
| sequence = gen_sequence(visorgpt_args, visorgpt_model, seq_prompt) | |
| torch.cuda.empty_cache() | |
| assert isinstance(sequence, list) | |
| try: | |
| cond_mask, cond_json = gen_cond_mask(sequence, ctn) | |
| if gen_type == 'key point': | |
| ori_sequence = cond_json[2]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'box': | |
| ori_sequence = cond_json[0]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'mask': | |
| ori_sequence = cond_json[1]['sequences'][0][0] + '[SEP]' | |
| except: | |
| cond_mask, cond_json = gen_cond_mask(sequence, not ctn) | |
| if gen_type == 'key point': | |
| ori_sequence = cond_json[2]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'box': | |
| ori_sequence = cond_json[0]['sequences'][0][0] + '[SEP]' | |
| elif gen_type == 'mask': | |
| ori_sequence = cond_json[1]['sequences'][0][0] + '[SEP]' | |
| ret_img = Image.fromarray(cond_mask) | |
| if not gen_type == 'mask': | |
| return { | |
| result_gallery: [ret_img], | |
| raw_sequence: gr.update(value=ori_sequence, visible=True), | |
| images_button: gr.update(visible=True), | |
| text_container: cond_json, | |
| sequence_container: ori_sequence | |
| } | |
| else: | |
| return { | |
| result_gallery: [ret_img], | |
| raw_sequence: gr.update(value=ori_sequence, visible=True), | |
| images_button: gr.update(visible=False), | |
| text_container: cond_json, | |
| sequence_container: ori_sequence | |
| } | |
| def generate_images(gen_type, | |
| num_samples, | |
| ddim_steps, | |
| object_prompt, | |
| seed, | |
| global_text, | |
| global_seq): | |
| if gen_type == 'key point': | |
| data = global_text[2]['keypoints'] | |
| idx = np.arange(len(data)) | |
| split_idx = list(np.array_split(idx, 1)[0]) | |
| for idx in tqdm(split_idx): | |
| item = data[idx] | |
| keypoint_list = [] | |
| for ins in item: | |
| kv = list(ins.items())[0] | |
| keypoint = (np.array(kv[1])).tolist() | |
| keypoint_list.append(keypoint) | |
| with torch.no_grad(): | |
| ret_img = control_infer(model=control_model, | |
| ddim_sampler=ddim_sampler, | |
| keypoint_list=keypoint_list, | |
| prompt=object_prompt.strip(), | |
| num_samples=num_samples, | |
| ddim_steps=ddim_steps, | |
| seed=seed) | |
| torch.cuda.empty_cache() | |
| elif gen_type == 'box': | |
| data = global_text[0]['bboxes'] | |
| with torch.no_grad(): | |
| ret_img = gligen_infer(model=g_model, | |
| autoencoder=g_autoencoder, | |
| text_encoder=g_text_encoder, | |
| diffusion=g_diffusion, | |
| config=g_config, | |
| grounding_tokenizer_input=g_grounding_tokenizer_input, | |
| context_prompt=object_prompt.strip(), | |
| bbox_lists=data, | |
| ddim_steps=ddim_steps, | |
| batch_size=num_samples, | |
| seed=seed) | |
| torch.cuda.empty_cache() | |
| if not gen_type == 'mask': | |
| return { | |
| result_gallery: ret_img, | |
| text_container: global_text, | |
| sequence_container: global_seq | |
| } | |
| else: | |
| return { | |
| raw_sequence: "sequence to mask is not supported yet :)", | |
| text_container: global_text, | |
| sequence_container: global_seq | |
| } | |
| def object_name_inbox_fn(gen_type): | |
| if gen_type == 'key point': | |
| return { | |
| object_name_inbox: gr.update(visible=False), | |
| data_type: gr.update(choices=['multiple instances']), | |
| images_button: gr.update(value='Synthesize images using ControlNet'), | |
| ddim_steps: gr.update(value=20), | |
| object_prompt: gr.update(placeholder='in suit'), | |
| num_instance: gr.update(visible=True, minimum=1, maximum=16, value=2, step=1), | |
| sequence_container: None | |
| } | |
| elif gen_type == 'box': | |
| return { | |
| object_name_inbox: gr.update(visible=True, value='person; frisbee'), | |
| data_type: gr.update(choices=['multiple instances', 'object centric']), | |
| images_button: gr.update(value='Synthesize images using GLIGEN'), | |
| ddim_steps: gr.update(value=50), | |
| object_prompt: gr.update(placeholder='man and frisbee'), | |
| num_instance: gr.update(visible=True, minimum=1, maximum=16, value=2, step=1), | |
| sequence_container: None | |
| } | |
| elif gen_type == 'mask': | |
| return { | |
| object_name_inbox: gr.update(visible=True, | |
| label="MS COCO categories to be generated (separated by semicolon)", value='bottle; cup'), | |
| data_type: gr.update(choices=['multiple instances']), | |
| images_button: gr.update(value='Synthesize images using GLIGEN'), | |
| ddim_steps: gr.update(value=50), | |
| object_prompt: gr.update(placeholder='bottle and cup'), | |
| num_instance: gr.update(visible=True, minimum=1, maximum=8, value=2, step=1), | |
| sequence_container: None | |
| } | |
| def instance_type_change_fn(data_type): | |
| if data_type == 'multiple instances': | |
| return { | |
| md_title: gr.update(visible=True), | |
| num_continuous_gen: gr.update(visible=True), | |
| continuous_btn: gr.update(visible=True), | |
| object_name_inbox: gr.update(label="MS COCO categories to be generated (separated by semicolon)", value='person; frisbee'), | |
| object_prompt: gr.update(placeholder='man and frisbee'), | |
| num_instance: gr.update(visible=True, minimum=1, maximum=16, value=2, step=1), | |
| } | |
| elif data_type == 'object centric': | |
| return { | |
| md_title: gr.update(visible=False), | |
| num_continuous_gen: gr.update(visible=False), | |
| continuous_btn: gr.update(visible=False), | |
| object_name_inbox: gr.update(label="ImageNet-1K categories to be generated", value='great white shark'), | |
| object_prompt: gr.update(placeholder='great white shark'), | |
| num_instance: gr.update(visible=False, value=1), | |
| } | |
| block = gr.Blocks() | |
| with block: | |
| text_container = gr.State() | |
| sequence_container = gr.State() | |
| gr.Markdown('<div align=center> <img src="file/visorgpt_title_all.jpg" width = "100%" height = "100%" /> </div>') | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### Params to generate sequences") | |
| gen_type = gr.inputs.Dropdown(choices=['key point', 'box', 'mask'], type='value', default='key point', label='Anotation Type') | |
| data_type = gr.inputs.Dropdown(choices=['multiple instances'], type='value', default='multiple instances', label='Data Type') | |
| instance_size = gr.inputs.Dropdown(choices=['small', 'medium', 'large'], type='value', default='large', label='Instance Size') | |
| num_instance = gr.Slider(label="Number of instances per image", minimum=1, maximum=16, value=2, step=1) | |
| object_name_inbox = gr.Textbox(label="MS COCO categories to be generated (separated by semicolon)", placeholder="person; frisbee", visible=False) | |
| sequence_button = gr.Button(value="Customize sequential output") | |
| md_title = gr.Markdown("### Continuous generation (Optional)") | |
| num_continuous_gen = gr.Slider(label="Add instances to the current scene", minimum=1, maximum=16, value=1, step=1) | |
| continuous_btn = gr.Button(value="Add instances to the current scene") | |
| gr.Markdown("### Params to synthesize images") | |
| object_prompt = gr.Textbox(label="Context Prompt", placeholder="in suit", visible=True) | |
| num_samples = gr.Slider(label="Batch Size", minimum=1, maximum=36, value=1, step=1) | |
| ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) | |
| seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True) | |
| images_button = gr.Button(value="Synthesize images using ControlNet", visible=False) | |
| with gr.Column(): | |
| raw_sequence = gr.Textbox(label="Raw Sequence", visible=False) | |
| result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto', preview=True) | |
| gen_type.change(object_name_inbox_fn, inputs=[gen_type], | |
| outputs=[object_name_inbox, data_type, images_button, ddim_steps, object_prompt, num_instance, sequence_container]) | |
| data_type.change(instance_type_change_fn, inputs=[data_type], | |
| outputs=[md_title, num_continuous_gen, continuous_btn, object_name_inbox, object_prompt, num_instance]) | |
| ips = [gen_type, data_type, instance_size, num_instance, object_name_inbox] | |
| sequence_button.click(fn=generate_sequence, inputs=ips, outputs=[result_gallery, raw_sequence, images_button, text_container, sequence_container]) | |
| ips = [gen_type, data_type, instance_size, num_instance, object_name_inbox, num_continuous_gen, sequence_container] | |
| continuous_btn.click(fn=add_contents, inputs=ips, outputs=[result_gallery, raw_sequence, images_button, text_container, sequence_container]) | |
| ips = [gen_type, num_samples, ddim_steps, object_prompt, seed, text_container, sequence_container] | |
| images_button.click(fn=generate_images, inputs=ips, outputs=[result_gallery, raw_sequence, text_container, sequence_container]) | |
| block.launch(server_name='0.0.0.0', server_port=10086, debug=False, share=False) | |