fixed model

inferencer.py

@@ -2,16 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-from typing import List, Dict, Tuple, Optional, Union, Any
-import matplotlib.pyplot as plt
+from typing import List, Optional, Union, Any
 
 from PIL import Image
 import torch
-import torch.nn.functional as F
-from torch import nn
-from torch.nn.attention.flex_attention import create_block_mask
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_utils import PreTrainedModel
 
 from data.data_utils import pil_img2rgb
 from modeling.bagel.qwen2_navit import NaiveCache
@@ -196,17 +190,17 @@ class InterleaveInferencer:
         ropes = gen_context['ropes']
 
         generation_input = self.model.prepare_start_tokens(kv_lens, ropes, self.new_token_ids)
-        unpacked_latent = self.model.generate_text(
+        for unpacked_latent in self.model.generate_text(
             past_key_values=past_key_values,
             max_length=max_length,
             do_sample=do_sample,
             temperature=temperature,
             end_token_id=self.new_token_ids['eos_token_id'],
             **generation_input,
-        )
-        output = self.tokenizer.decode(unpacked_latent[:,0])
-        output = output.split('<|im_end|>')[0].split('<|im_start|>')[1]
-        return output
+        ):
+            output = self.tokenizer.decode(unpacked_latent)
+            if output != "<|im_end|>":
+                yield output
         
     @torch.no_grad()
     def interleave_inference(
@@ -214,10 +208,11 @@ class InterleaveInferencer:
         input_lists: List[Union[str, Image.Image]],
         think=False,
         understanding_output=False,
-
+        # for gen_text
         max_think_token_n=1000,
         do_sample=False,
         text_temperature=0.3,
+        # for gen_image
         cfg_text_scale=3.0,
         cfg_img_scale=1.5,
         cfg_interval=[0.4, 1.0],
@@ -225,23 +220,20 @@ class InterleaveInferencer:
         num_timesteps=50,
         cfg_renorm_min=0.0,
         cfg_renorm_type="global",
-        image_shapes=(1024, 1024),
-    ) -> List[Union[str, Image.Image]]:
-
-        output_list = []
+        image_shapes=(1024, 1024), # Default, can be overridden by actual input image
+    ):
         gen_context = self.init_gen_context()
         cfg_text_context = deepcopy(gen_context)
         cfg_img_context = deepcopy(gen_context)
 
         with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
+
             if think:
-                if understanding_output:
-                    system_prompt = VLM_THINK_SYSTEM_PROMPT 
-                else:
-                    system_prompt = GEN_THINK_SYSTEM_PROMPT
+                system_prompt = VLM_THINK_SYSTEM_PROMPT if understanding_output else GEN_THINK_SYSTEM_PROMPT
                 gen_context = self.update_context_text(system_prompt, gen_context)
+                cfg_text_context = self.update_context_text(system_prompt, cfg_text_context)
                 cfg_img_context = self.update_context_text(system_prompt, cfg_img_context)
-
+    
             for input_term in input_lists:
                 if isinstance(input_term, str):
                     cfg_text_context = deepcopy(gen_context)
@@ -251,29 +243,29 @@ class InterleaveInferencer:
                 elif isinstance(input_term, Image.Image):
                     input_term = self.vae_transform.resize_transform(pil_img2rgb(input_term))
                     gen_context = self.update_context_image(input_term, gen_context, vae=not understanding_output)
-
                     image_shapes = input_term.size[::-1]
                     cfg_text_context = deepcopy(gen_context)
 
                 else:
                     raise ValueError(f"Unsupported input type: {type(input_term)}")
-
-            if understanding_output:
-                gen_text = self.gen_text(gen_context, do_sample=do_sample, temperature=text_temperature, max_length=max_think_token_n)
-                output_list.append(gen_text)
-
-            else:
+    
+            if understanding_output: # Generate text
+                yield from self.gen_text(gen_context, max_length=max_think_token_n, do_sample=do_sample, temperature=text_temperature)
+            else: # Generate image
                 if think:
-                    gen_text = self.gen_text(gen_context, do_sample=do_sample, temperature=text_temperature, max_length=max_think_token_n)
-                    gen_context = self.update_context_text(gen_text, gen_context)
-                    output_list.append(gen_text)
-
+                    thought_text_parts = []
+                    for part in self.gen_text(gen_context, max_length=max_think_token_n, do_sample=do_sample, temperature=text_temperature):
+                        yield part # Stream the thought
+                        thought_text_parts.append(part)
+                    full_thought_text = "".join(thought_text_parts)
+                    if full_thought_text: # Only update if thought was generated
+                        gen_context = self.update_context_text(full_thought_text, gen_context)
+                
                 img = self.gen_image(
-                    image_shapes, 
-                    gen_context, 
+                    image_shape=image_shapes, 
+                    gen_context=gen_context, 
                     cfg_text_precontext=cfg_text_context, 
                     cfg_img_precontext=cfg_img_context,
-
                     cfg_text_scale=cfg_text_scale, 
                     cfg_img_scale=cfg_img_scale, 
                     cfg_interval=cfg_interval, 
@@ -282,34 +274,21 @@ class InterleaveInferencer:
                     cfg_renorm_min=cfg_renorm_min,
                     cfg_renorm_type=cfg_renorm_type,
                 )
+                yield img
 
-                output_list.append(img)
-
-        return output_list
-    
     def __call__(
         self, 
         image: Optional[Image.Image] = None, 
         text: Optional[str] = None, 
-        **kargs
-    ) -> Dict[str, Any]:
-        output_dict = {'image': None, 'text': None}
-
-        if image is None and text is None:
-            print('Please provide at least one input: either an image or text.')
-            return output_dict
-
+        **kargs 
+    ) -> Any:
         input_list = []
         if image is not None:
             input_list.append(image)
         if text is not None:
             input_list.append(text)
+        
+        if not input_list and not kargs.get('force_empty_input', False): # allow forcing for special cases if needed
+            return
 
-        output_list = self.interleave_inference(input_list, **kargs)
-
-        for i in output_list:
-            if isinstance(i, Image.Image):
-                output_dict['image'] = i
-            elif isinstance(i, str):
-                output_dict['text'] = i
-        return output_dict
+        yield from self.interleave_inference(input_list, **kargs)

modeling/bagel/bagel.py

@@ -890,16 +890,21 @@ class Bagel(PreTrainedModel):
         temperature: float = 1.0,
         end_token_id: int = None,
     ):
+        """
+        Generates text token by token in a streaming fashion.
+
+        This function is a generator that yields one token at a time. It replicates
+        the behavior of the original batch generation function, including the handling
+        of start tokens and the end-of-sequence token.
+        """
         step = 0
-        generated_sequence = []
         curr_tokens = packed_start_tokens
         while step < max_length:
-            generated_sequence.append(curr_tokens)
             packed_text_embedding = self.language_model.model.embed_tokens(curr_tokens)
             query_lens = torch.ones_like(curr_tokens)
             packed_query_indexes = torch.cumsum(key_values_lens, dim=0) + torch.arange(
-                0, len(key_values_lens), 
-                device=key_values_lens.device, 
+                0, len(key_values_lens),
+                device=key_values_lens.device,
                 dtype=key_values_lens.dtype
             )
 
@@ -944,12 +949,11 @@ class Bagel(PreTrainedModel):
             packed_query_position_ids = packed_query_position_ids + 1
             step += 1
 
+            yield curr_tokens # Yield each token as it's generated
+
             if end_token_id is not None and curr_tokens[0] == end_token_id: # only support batch=1
                 break
 
-        output_device = generated_sequence[0].device
-        return torch.stack([i.to(output_device) for i in generated_sequence], dim=0)
-
     # for evaluation
     @torch.no_grad()
     def chat(
@@ -1012,15 +1016,13 @@ class Bagel(PreTrainedModel):
             if torch.is_tensor(v):
                 generation_input[k] = v.to(device)
         with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
-            unpacked_latent = self.generate_text(
+            for unpacked_latent in self.generate_text(
                 past_key_values=past_key_values,
                 max_length=max_length,
                 do_sample=do_sample,
                 temperature=temperature,
                 end_token_id=new_token_ids['eos_token_id'],
                 **generation_input,
-            )
-        output = tokenizer.decode(unpacked_latent[:,0])
-        output = output.split('<|im_end|>')[0].split('<|im_start|>')[1]
-
-        return output
+            ):
+                output = tokenizer.decode(unpacked_latent[:,0])
+                yield output