Upload moondream.py

moondream.py

@@ -62,7 +62,6 @@ class Moondream(PreTrainedModel):
     def get_input_embeddings(self):
         return self.text_model.get_input_embeddings()
 
-
     def generate(
         self,
         image_embeds,
@@ -85,17 +84,9 @@ class Moondream(PreTrainedModel):
             output_ids = self.text_model.generate(
                 inputs_embeds=inputs_embeds, streamer=streamer, **generate_config
             )
+            print("FINISHED")
 
-            # Pull the generated text from the streamer, and update the model output.
-            model_output = ""
-            for new_text in streamer:
-                model_output += new_text
-                yield new_text
-            return model_output
-
-        #return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-
-
+        return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
 
     def answer_question(
         self,
@@ -186,4 +177,4 @@ class Moondream(PreTrainedModel):
         return [
             x.strip()
             for x in tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-        ]
+        ]