Upload moondream.py

moondream.py

@@ -1,7 +1,7 @@
 import torch
 from .vision_encoder import VisionEncoder
 from .configuration_moondream import MoondreamConfig
-from transformers import PreTrainedModel, TextStreamer
+from transformers import PreTrainedModel, TextIteratorStreamer
 
 from .modeling_phi import PhiForCausalLM
 from .configuration_moondream import PhiConfig
@@ -62,6 +62,7 @@ class Moondream(PreTrainedModel):
     def get_input_embeddings(self):
         return self.text_model.get_input_embeddings()
 
+
     def generate(
         self,
         image_embeds,
@@ -80,13 +81,15 @@ class Moondream(PreTrainedModel):
 
         with torch.no_grad():
             inputs_embeds = self.input_embeds(prompt, image_embeds, tokenizer)
-            streamer = TextStreamer(tokenizer)
+            streamer = TextIteratorStreamer(tokenizer)
             output_ids = self.text_model.generate(
                 inputs_embeds=inputs_embeds, streamer=streamer, **generate_config
             )
 
         return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
 
+
+
     def answer_question(
         self,
         image_embeds,