Upload Moondream

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7eb78782910cad8e8ba5a90146a640ef48b420524d8e946ecb65519834145acc
+oid sha256:493ac8972766b8e4b9005bfab11454b93aab4987b44a01debebec3fa96773105
 size 3715037856

moondream.py

@@ -9,13 +9,16 @@ from .configuration_moondream import PhiConfig
 
 class Moondream(PreTrainedModel):
     config_class = MoondreamConfig
+    _supports_flash_attn_2 = True
 
     def __init__(self, config):
         super().__init__(config)
         self.vision_encoder = VisionEncoder()
 
         if type(config.phi_config) == dict:
-            phi_config = PhiConfig(**config.phi_config)
+            phi_config = PhiConfig(
+                **config.phi_config, attn_implementation=config._attn_implementation
+            )
         else:
             phi_config = config.phi_config
         self.text_model = PhiForCausalLM(phi_config)
@@ -94,7 +97,7 @@ class Moondream(PreTrainedModel):
             prompt,
             eos_text="<END>",
             tokenizer=tokenizer,
-            max_new_tokens=256,
+            max_new_tokens=512,
             **kwargs,
         )[0]
         cleaned_answer = re.sub("<$|<END$", "", answer).strip()
@@ -104,3 +107,71 @@ class Moondream(PreTrainedModel):
             result_queue.put(cleaned_answer)
         else:
             return cleaned_answer
+
+    def batch_answer(
+        self,
+        images,
+        prompts,
+        tokenizer,
+        **kwargs,
+    ):
+        eos_tokens = tokenizer("<END>", add_special_tokens=False)[0].ids
+
+        image_embeds = self.encode_image(images)
+
+        templated_prompts = [
+            f"<image>\n\nQuestion: {prompt}\n\nAnswer: " for prompt in prompts
+        ]
+        prompt_embs = [
+            self.input_embeds(prompt, image_embed.unsqueeze(0), tokenizer)[0]
+            for prompt, image_embed in zip(templated_prompts, image_embeds)
+        ]
+
+        bos_emb = prompt_embs[0][0]
+        max_len = max([p.shape[0] for p in prompt_embs])
+
+        inputs_embeds = torch.cat(
+            [
+                torch.cat([bos_emb.repeat(max_len - p.shape[0], 1), p]).unsqueeze(0)
+                for p in prompt_embs
+            ],
+            dim=0,
+        )
+        attention_mask = torch.cat(
+            [
+                torch.cat(
+                    [
+                        torch.zeros(
+                            1,
+                            max_len - p.shape[0],
+                            device=self.device,
+                            dtype=torch.long,
+                        ),
+                        torch.ones(1, p.shape[0], device=self.device, dtype=torch.long),
+                    ],
+                    dim=1,
+                )
+                for p in prompt_embs
+            ],
+            dim=0,
+        )
+
+        generate_config = {
+            "eos_token_id": eos_tokens,
+            "bos_token_id": tokenizer.bos_token_id,
+            "pad_token_id": tokenizer.eos_token_id,
+            "max_new_tokens": 512,
+            **kwargs,
+        }
+
+        with torch.no_grad():
+            output_ids = self.text_model.generate(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                **generate_config,
+            )
+
+        return [
+            re.sub("<$|<END$", "", x).strip()
+            for x in tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        ]

vision_encoder.py

@@ -121,13 +121,15 @@ class VisionEncoder(nn.Module):
     def dtype(self):
         return self.projection.mlp.fc1.weight.dtype
 
-    def __call__(self, image: Image) -> torch.Tensor:
+    def __call__(self, images) -> torch.Tensor:
+        if not isinstance(images, list):
+            images = [images]
+
         with torch.no_grad():
-            x = (
-                self.preprocess(image.convert("RGB"))
-                .unsqueeze(0)
-                .to(self.device, dtype=self.dtype)
-            )
+            x = torch.stack(
+                [self.preprocess(image.convert("RGB")) for image in images]
+            ).to(self.device, dtype=self.dtype)
+
             x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14)
 
             x = self.encoder(x)