demo

app.py

@@ -37,6 +37,7 @@ IMG_TOKEN = '<img_{:05d}>'
 IMG_FLAG = '<image>'
 num_img_in_tokens = 64
 num_img_out_tokens = 64
+instruction_prompt = '{instruction}'
 
 resolution_grids = ['1x1', '1x2', '1x3', '1x4', '1x5', '1x6', '1x10', '2x1', '3x1', '4x1', '5x1', '6x1', '10x1', '2x2',
                     '2x3', '3x2', '2x4', '4x2']
@@ -119,6 +120,7 @@ class LLMService:
         self.agent = hydra.utils.instantiate(agent_cfg, llm=llm)
 
         self.agent.eval().to(self.llm_device, dtype=self.dtype)
+        self.agent.llm.base_model.model.use_kv_cache_head = False
         print('Init agent mdoel Done')
 
         noise_scheduler = EulerDiscreteScheduler.from_pretrained(args.diffusion_path, subfolder="scheduler")
@@ -166,10 +168,13 @@ service = LLMService(args)
 
 
 @spaces.GPU
-def generate(text_list, image_list, max_new_tokens):
+def generate(text_list, image_list, image_embed_list, max_new_tokens):
     with torch.no_grad():
+        print('text_list: {}'.format(text_list))
         text_list = text_list.split(IMG_FLAG)
+        text_list = [text_list[0]] + ["[INST]"+item for item in text_list[1:-1]] + [text_list[-1]]
         top_p = 0.5
+        window_size = 8
         assert len(text_list) == len(image_list) + 1
 
         image_tokens = BOI_TOKEN + ''.join(
@@ -181,31 +186,12 @@ def generate(text_list, image_list, max_new_tokens):
             embeds_cmp_mask = []
             embeds_gen_mask = []
 
-            if service.multi_resolution:
-                patch_pos = []
-                image_patch_length = []
-                image_size_list = []
-
             for idx, image_item in enumerate(image_list):
                 if isinstance(image_item, str):
                     image = decode_image(image_item)
                     print('after decode image size:', image.size)
                     input_images.append(image)
 
-                    # if service.multi_resolution:
-                    #     image_size_list.append(image.size)
-                    #     print('image size:', image.size)
-                    #     image_tensor, patch_pos_tensor = process_anyres_image(image, service.image_transform,
-                    #                                                           service.grid_pinpoints,
-                    #                                                           service.base_resolution)
-                    #     image_tensor_list.append(image_tensor)
-                    #     patch_pos.append(patch_pos_tensor)
-                    #     image_patch_length.append(image_tensor.shape[0])
-                    #     print('image_patch_length', image_patch_length)
-                    #     embeds_cmp_mask.extend([True] * image_tensor.shape[0])
-                    #     embeds_gen_mask.extend([False] * image_tensor.shape[0])
-                    #
-                    # else:
                     image_tensor = service.image_transform(image)
                     image_tensor_list.append(image_tensor)
                     embeds_cmp_mask.append(True)
@@ -213,23 +199,13 @@ def generate(text_list, image_list, max_new_tokens):
                 else:
                     raise ValueError
 
-            if service.multi_resolution:
-                pixel_values = torch.cat(image_tensor_list).to(service.vit_sd_device, dtype=service.dtype)
-                patch_position = torch.cat(patch_pos, dim=0)
-
-                image_tokens_list = []
-                for patch_length in image_patch_length:
-                    image_tokens = ''
-                    for _ in range(patch_length - 1):
-                        image_tokens += BOP_TOKEN + ''.join(
-                            IMG_TOKEN.format(int(item)) for item in range(num_img_in_tokens)) + EOP_TOKEN
-                    image_tokens += BOI_TOKEN + ''.join(
-                        IMG_TOKEN.format(int(item)) for item in range(num_img_in_tokens)) + EOI_TOKEN
-                    image_tokens_list.append(image_tokens)
-            else:
-                pixel_values = torch.stack(image_tensor_list).to(service.vit_sd_device, dtype=service.dtype)
-
-            image_embeds = service.visual_encoder(pixel_values)
+            # pixel_values = torch.stack(image_tensor_list).to(service.vit_sd_device, dtype=service.dtype)
+            # 
+            # image_embeds = service.visual_encoder(pixel_values)
+            # image_embeds = image_embeds.to(service.llm_device)
+            print(image_embed_list)
+            image_embed_list = [t.squeeze(0) for t in image_embed_list]
+            image_embeds = torch.stack(image_embed_list, dim=0)
             image_embeds = image_embeds.to(service.llm_device)
 
             embeds_cmp_mask = torch.tensor(embeds_cmp_mask, dtype=torch.bool).to(service.llm_device)
@@ -241,10 +217,21 @@ def generate(text_list, image_list, max_new_tokens):
             embeds_cmp_mask = None
             embeds_gen_mask = None
 
+
         input_text = image_tokens.join(text_list)
 
-        print('input_text:', input_text)
+        print('input_text fed to LLM:', input_text)
         input_ids = service.tokenizer.encode(input_text, add_special_tokens=False)
+
+        while image_embeds.shape[0] > window_size:
+            eoi_prompt_idx = input_text.index(EOI_TOKEN)
+            input_text = input_text[eoi_prompt_idx + len(EOI_TOKEN) + len('[INST]'):]
+            image_embeds = image_embeds[1:]
+            input_ids = service.tokenizer.encode(input_text, add_special_tokens=False)
+
+        if image_embeds is not None:
+            embeds_cmp_mask = torch.tensor([True] * image_embeds.shape[0]).to(service.llm_device, dtype=torch.bool)
+
         input_ids = [service.tokenizer.bos_token_id] + input_ids
 
         input_ids = torch.tensor(input_ids).to(service.llm_device, dtype=torch.long)
@@ -262,7 +249,10 @@ def generate(text_list, image_list, max_new_tokens):
         ids_gen_mask = ids_gen_mask.unsqueeze(0)
 
         error_msg = []
-
+        print('image_embeds_shape: ' + str(image_embeds.shape))
+        print('image_embeds: {}'.format(image_embeds))
+        print('input_ids: ' + str(input_ids))
+        print('ids_cmp_mask: ' + str(ids_cmp_mask))
         output = service.agent.generate(
             tokenizer=service.tokenizer,
             input_ids=input_ids,
@@ -278,7 +268,6 @@ def generate(text_list, image_list, max_new_tokens):
 
         gen_imgs_base64_list = []
         generated_text = output['text']
-        generated_text = generated_text.replace(EOI_TOKEN, IMG_FLAG).replace(service.tokenizer.eos_token, '')
 
         torch.cuda.empty_cache()
 
@@ -294,6 +283,7 @@ def generate(text_list, image_list, max_new_tokens):
             for img_idx in range(output['num_gen_imgs']):
                 img_feat = img_gen_feat[img_idx:img_idx + 1]
                 generated_image = service.sd_adapter.generate(image_embeds=img_feat, num_inference_steps=50)[0]
+                gen_imgs_base64_list.append(generated_image)
 
             # a = time.time()
             # service.sd_adapter = service.sd_adapter.cpu()
@@ -301,19 +291,19 @@ def generate(text_list, image_list, max_new_tokens):
             # service.agent = service.agent.to(service.vit_sd_device, dtype=service.dtype)
             # print("Loading finished: ", time.time() - a)
 
-        print(input_text + generated_text)
-        return {'text': generated_text, 'images': gen_imgs_base64_list, 'error_msg': error_msg}
+        print('[func generate inout+output]: {}'.format(input_text + generated_text))
+        return {'text': generated_text, 'images': gen_imgs_base64_list, 'image_embeds': img_feat.detach().clone(), 'error_msg': error_msg}
 
 
-def http_bot(dialog_state, input_state, max_new_tokens, max_turns,
+def http_bot(dialog_state, input_state, max_new_tokens, max_length,
              request: gr.Request):
     print('input_state:', input_state)
-
-    if len(dialog_state.messages) == 0 or dialog_state.messages[-1]['role'] != dialog_state.roles[0] or len(
+    print(dialog_state.messages)
+    if len(dialog_state.messages) == 0 or len(
             dialog_state.messages[-1]['message']['text'].strip(' ?.;!/')) == 0:
         return (dialog_state, input_state, dialog_state.to_gradio_chatbot()) + (no_change_btn,) * 4
 
-    if len(dialog_state.messages) > max_turns * 2:
+    if len(dialog_state.messages) >= max_length:
         output_state = init_input_state()
         output_state['text'] = 'Error: History exceeds maximum rounds, please clear history and restart.'
         dialog_state.messages.append({'role': dialog_state.roles[1], 'message': output_state})
@@ -322,22 +312,40 @@ def http_bot(dialog_state, input_state, max_new_tokens, max_turns,
 
     prompt = dialog_state.get_prompt()
     text = prompt['text']
+    print('text from http_bot: {}'.format(text))
     max_new_tokens = int(max_new_tokens)
     images = prompt['images']
+    image_embeds = prompt['image_embeds']
+
+    results = generate(text, images, image_embeds, max_new_tokens)
+    generated_text = results['text']
+    pattern = r' <img_000\d{2}>'
+    # Replace all occurrences of the pattern with the replacement text
+    generated_text = re.sub(pattern, '', generated_text)
+
+    generated_text = generated_text.replace(' '+service.tokenizer.eos_token, '')\
+        .replace('[INST]', '').replace(' '+BOI_TOKEN, '').replace(' '+EOI_TOKEN, IMG_FLAG)
+
+    results['text'] = generated_text
 
-    results = generate(text, images, max_new_tokens)
     print('response: ', {'text': results['text'], 'error_msg': results['error_msg']})
 
     output_state = init_input_state()
     image_dir = get_conv_image_dir()
     output_state['text'] = results['text']
+    output_state['image_embeds'].append(results['image_embeds'])
 
     for image_base64 in results['images']:
         if image_base64 == '':
             image_path = ''
         else:
-            image = decode_image(image_base64)
-            image = image.convert('RGB')
+            if isinstance(image_base64, Image.Image):
+                print('generated image is in Image.Image')
+                image = image_base64
+            else:
+                print('generated image is in Image_base64')
+                image = decode_image(image_base64)
+                image = image.convert('RGB')
             image_path = get_image_name(image=image, image_dir=image_dir)
             if not os.path.exists(image_path):
                 image.save(image_path)
@@ -354,8 +362,8 @@ def http_bot(dialog_state, input_state, max_new_tokens, max_turns,
 IMG_FLAG = '<image>'
 LOGDIR = 'log'
 
-logger = build_logger("gradio_seed_x", LOGDIR)
-headers = {"User-Agent": "SEED-X Client"}
+logger = build_logger("gradio_seed_story", LOGDIR)
+headers = {"User-Agent": "SEED-Story Client"}
 
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
@@ -436,10 +444,16 @@ def center_crop_image(image, max_aspect_ratio=1.5):
 
 def vote_last_response(state, vote_type, request: gr.Request):
     with open(get_conv_log_filename(), "a") as fout:
+        print(state)
+        print(state.dict())
+        dic = state.dict()
+        for i in range(len(dic['messages'])):
+            dic['messages'][i]['message'].pop('image_embeds')
+        print(dic)
         data = {
             "tstamp": round(time.time(), 4),
             "type": vote_type,
-            "state": state.dict(),
+            "state": dic,
             "ip": request.client.host,
         }
         fout.write(json.dumps(data) + "\n")
@@ -475,7 +489,7 @@ def clear_history(request: gr.Request):
 
 
 def init_input_state():
-    return {'images': [], 'text': ''}
+    return {'images': [], 'text': '', 'image_embeds': []}
 
 
 def add_text(dialog_state, input_state, text, request: gr.Request):
@@ -509,13 +523,17 @@ def add_image(dialog_state, input_state, image, request: gr.Request):
 
     print('image size:', image.size)
 
-    image = center_crop_image(image, max_aspect_ratio=10)
+    # image = center_crop_image(image, max_aspect_ratio=10)
 
     image_dir = get_conv_image_dir()
     image_path = get_image_name(image=image, image_dir=image_dir)
     if not os.path.exists(image_path):
         image.save(image_path)
     input_state['images'].append(image_path)
+    image_tensor = service.image_transform(image).unsqueeze(0).to(service.llm_device, dtype=service.dtype)
+    image_embeds = service.visual_encoder(image_tensor).detach().clone()
+    image_embeds = image_embeds.to(service.llm_device)
+    input_state['image_embeds'].append(image_embeds)
     input_state['text'] += IMG_FLAG
 
     if len(dialog_state.messages) > 0 and dialog_state.messages[-1]['role'] == dialog_state.roles[0]:
@@ -548,14 +566,13 @@ title = ("""
 # SEED-Story
 [[Paper]](https://arxiv.org/abs/2407.08683) [[Code]](https://github.com/TencentARC/SEED-Story)
 
-Demo of a multimodal story generation model SEED-Story-George. It is trained on StoryStream-Curious George subset.
+Demo of the multimodal story generation model SEED-Story-George. It is trained on StoryStream-Curious George subset.
 SEED-Story is a MLLM capable of generating multimodal long stories consisting of rich and coherent narrative texts, along with images that are consistent in characters and style.
 
 ## Tips:
 * Check out the conversation examples (at the bottom) for inspiration.
-* You can adjust "Max History Rounds" to try a conversation with up to **three rounds due to insufficient GPU memory**. For more turns, you can download our checkpoints from GitHub and deploy them locally for inference.
-* Our demo supports a mix of images and texts as input. You can freely upload an image or enter text, and then click on "Add Image/Text". You can repeat the former step multiple times, and click on "Submit" for model inference at last.
-
+* Our demo requires a mix of an image and a starting sentence as input. You can freely upload an image or enter text, and then click on "Submit". Then, The model generates the next story image and text.
+* You can click on "Continue Generation" to make the model generate a next story image and text based on all previous story boards.
 * SEED-Story was trained with English-only data. It may process with other languages due to the inherent capabilities from LLaMA, but might not stable.
 """)
 
@@ -577,7 +594,7 @@ img:before {
   position: absolute;
   top: -10px;
   left: 0;
-  height: calc(100% + 10px);
+  height: auto;
   width: 100%;
   background-color: rgb(230, 230, 230);
   border: 2px dotted rgb(200, 200, 200);
@@ -601,29 +618,10 @@ img:after {
 
 if __name__ == '__main__':
     examples_mix = [
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/bank.png?raw=true',
-         'Can I conntect with an advisor on Sunday?'],
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/ground.png?raw=true',
-         'Is there anything in the image that can protect me from catching the flu virus when I go out? Show me the location.'],
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/arrow.jpg?raw=true',
-         'What is the object pointed by the red arrow?'],
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/shanghai.png?raw=true',
-         'Where was this image taken? Explain your answer.'],
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/GPT4.png?raw=true',
-         'How long does it take to make GPT-4 safer?'],
-        ['https://github.com/AILab-CVC/SEED-X/blob/main/demos/twitter.png?raw=true',
-         'Please provide a comprehensive description of this image.'],
-    ]
-    examples_text = [
-        ['I want to build a two story cabin in the woods, with many commanding windows. Can you show me a picture?'],
-        ['Use your imagination to design a concept image for Artificial General Intelligence (AGI). Show me an image.'],
-        [
-            'Can you design an illustration for “The Three-Body Problem” to depict a scene from the novel? Show me a picture.'],
-        [
-            'My four year old son loves toy trains. Can you design a fancy birthday cake for him? Please generate a picture.'],
-        [
-            'Generate an image of a portrait of young nordic girl, age 25, freckled skin, neck tatoo, blue eyes 35mm lens, photography, ultra details.'],
-        ['Generate an impressionist painting of an astronaut in a jungle.']
+        ['https://github.com/TencentARC/SEED-Story/blob/master/assets/demo_examples/2.jpg?raw=true',
+         'One day, George, the curious brown monkey, decided to explore a new room. He peeked out from behind a dresser, looking both curious and cautious. The dresser had three drawers, each with a round handle. An electrical outlet was visible on the wall.'],
+        ['https://github.com/TencentARC/SEED-Story/blob/master/assets/demo_examples/4.jpg?raw=true',
+         'In the bustling city, a beautiful blue and yellow bird took flight, soaring high above the buildings. Among the clouds, a heart-shaped formation appeared, as if nature was sending a love note to the world below. Other birds joined, their silhouettes dancing in the distance.'],
     ]
     with gr.Blocks(css=css) as demo:
         gr.Markdown(title)
@@ -640,10 +638,10 @@ if __name__ == '__main__':
                                       elem_id='textbox',
                                       placeholder="Enter text and image, and press submit,", container=False)
                 with gr.Row():
-                    add_image_btn = gr.Button("Add Image")
-                    add_text_btn = gr.Button("Add Text")
-
+                    # add_image_btn = gr.Button("Add Image")
+                    # add_text_btn = gr.Button("Add Text")
                     submit_btn = gr.Button("Submit")
+                    continue_btn = gr.Button("Continue Generation")
 
                 with gr.Row():
                     max_new_tokens = gr.Slider(minimum=64,
@@ -652,14 +650,11 @@ if __name__ == '__main__':
                                                step=64,
                                                interactive=True,
                                                label="Max Output Tokens")
-                    max_turns = gr.Slider(minimum=1, maximum=3, value=3, step=1, interactive=True,
-                                          label="Max History Rounds")
-                    force_img_gen = gr.Radio(choices=[True, False], value=False, label='Force Image Generation')
-                    force_bbox = gr.Radio(choices=[True, False], value=False, label='Force Bounding Box')
-                    force_polish = gr.Radio(choices=[True, False], value=True, label='Force Polishing Generated Image')
+                    max_length = gr.Slider(minimum=1, maximum=30, value=10, step=1, interactive=True,
+                                          label="Max Story Length")
 
             with gr.Column(scale=7):
-                chatbot = gr.Chatbot(elem_id='chatbot', label="SEED-X-I", height=700)
+                chatbot = gr.Chatbot(elem_id='chatbot', label="SEED-Story", height=700)
                 with gr.Row():
                     upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
                     downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
@@ -667,10 +662,8 @@ if __name__ == '__main__':
                     clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
 
         with gr.Row():
-            with gr.Column(scale=0.7):
+            with gr.Column(scale=1.0):
                 gr.Examples(examples=examples_mix, label='Input examples', inputs=[image, text], cache_examples=False)
-            with gr.Column(scale=0.3):
-                gr.Examples(examples=examples_text, label='Input examples', inputs=[text], cache_examples=False)
 
         # Register listeners
         btn_list = [upvote_btn, downvote_btn, regenerate_btn, clear_btn]
@@ -678,20 +671,25 @@ if __name__ == '__main__':
         downvote_btn.click(downvote_last_response, [dialog_state], [upvote_btn, downvote_btn])
 
         regenerate_btn.click(regenerate, [dialog_state], [dialog_state, chatbot] + btn_list).then(
-            http_bot, [dialog_state, input_state, max_new_tokens, max_turns, force_img_gen, force_bbox, force_polish],
+            http_bot, [dialog_state, input_state, max_new_tokens, max_length],
             [dialog_state, input_state, chatbot] + btn_list)
-        add_image_btn.click(add_image, [dialog_state, input_state, image],
-                            [dialog_state, input_state, image, chatbot] + btn_list)
-
-        add_text_btn.click(add_text, [dialog_state, input_state, text],
-                           [dialog_state, input_state, text, chatbot] + btn_list)
+        # add_image_btn.click(add_image, [dialog_state, input_state, image],
+        #                     [dialog_state, input_state, image, chatbot] + btn_list)
+        #
+        # add_text_btn.click(add_text, [dialog_state, input_state, text],
+        #                    [dialog_state, input_state, text, chatbot] + btn_list)
 
         submit_btn.click(
-            add_image, [dialog_state, input_state, image], [dialog_state, input_state, image, chatbot] + btn_list).then(
             add_text, [dialog_state, input_state, text],
             [dialog_state, input_state, text, chatbot, upvote_btn, downvote_btn, regenerate_btn, clear_btn]).then(
+            add_image, [dialog_state, input_state, image],
+            [dialog_state, input_state, image, chatbot] + btn_list).then(
+            http_bot,
+            [dialog_state, input_state, max_new_tokens, max_length],
+            [dialog_state, input_state, chatbot] + btn_list)
+        continue_btn.click(
             http_bot,
-            [dialog_state, input_state, max_new_tokens, max_turns, force_img_gen, force_bbox, force_polish],
+            [dialog_state, input_state, max_new_tokens, max_length],
             [dialog_state, input_state, chatbot] + btn_list)
         clear_btn.click(clear_history, None, [dialog_state, input_state, chatbot] + btn_list)
 

configs/visual_tokenizer/qwen_vitg_448.yaml

@@ -7,4 +7,4 @@ mlp_ratio: 4.9231
 output_dim: 4096
 patch_size: 14
 width: 1664
-pretrained_model_path: /dataset/syang/pretrained/qwen_vit_G.pt
+pretrained_model_path: pretrained/qwen_vit_G.pt

conversation.py

@@ -49,7 +49,8 @@ class Conversation:
     skip_next: bool = False
 
     def get_prompt(self):
-        messages = copy.deepcopy(self.messages)
+        messages = self.messages
+        # messages = copy.deepcopy(self.messages)
         if self.sep_style == SeparatorStyle.SINGLE:
             if self.system is None or self.system == '':
                 text = ''
@@ -65,28 +66,28 @@ class Conversation:
 
             text += self.roles[1] + ":"
         elif self.sep_style == SeparatorStyle.LLAMA_2:
-            b_token = "[INST] "
-            e_token = " [/INST]"
+            # b_token = "[INST] "
+            # e_token = " [/INST]"
             if self.system is None or self.system == '':
                 text = ''
             else:
                 text = f"<<SYS>>\n{self.system}\n<</SYS>>\n\n"
             images = []
+            image_embeds = []
             for idx, message in enumerate(messages):
-                # text += message['role'] + ": " + message['message']['text'] + self.sep
-                if idx % 2 == 0:
-                    text += b_token + message['message']['text'] + e_token + self.sep
-                else:
-                    text += message['message']['text'] + self.sep
+                text += message['message']['text']
 
                 for image_path in message['message']['images']:
-                    image = Image.open(image_path)
+                    image = Image.open(image_path).convert('RGB')
                     image_base64 = encode_image(image)
                     images.append(image_base64)
+                image_embeds.extend(message['message']['image_embeds'])
+
+
         else:
             raise NotImplementedError
 
-        return {'text': text, 'images': images}
+        return {'text': text, 'images': images, 'image_embeds': image_embeds}
 
     # def update_image_ids(self, images_ids):
     #     image_count = 0
@@ -106,6 +107,7 @@ class Conversation:
         for i, single_turn in enumerate(self.messages[self.offset:]):
             single_turn = single_turn['message']
             text_list = single_turn['text'].split(IMG_FLAG)
+            print(text_list, len(single_turn['images']))
             assert len(text_list) == len(single_turn['images']) + 1, print(text_list, len(single_turn['images']))
             message = ''
             for image_idx in range(len(single_turn['images'])):

pretrained/seed_story/george_sft/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:46db6f1beb672085204ca9f7d542f6b62063cbe9970933ca702bccc72f00a4f6
+oid sha256:c7e46794a2aab38f3f59484a4f4bb4c839217ef17c4329977b0a11839f462b94
 size 14709979626