Model update

README.md

@@ -209,7 +209,7 @@ python test_jvlm.py -i assets/the_persistence_of_memory.jpg -p "What's in this i
 python test_jvlm.py -i https://picsum.photos/id/1025/800/600.jpg -p "Describe this image"
 
 # Multiple images (local and remote)
-python test_jvlm.py -i https://picsum.photos/id/1015/800/600.jpg -i https://picsum.photos/id/1016/800/600.jpg -i https://picsum.photos/id/1021/800/600.jpg -p "What is the difference between these pictures?"
+python test_jvlm.py -i https://picsum.photos/id/1015/800/600.jpg -i https://picsum.photos/id/1016/800/600.jpg -i https://picsum.photos/id/1021/800/600.jpg -p "Describe these images"
 
 # Text only input
 python test_jvlm.py -p "How many planets are in our solar system?"
@@ -302,7 +302,7 @@ model = AutoModelForCausalLM.from_pretrained(
 #
 # model = AutoModelForCausalLM.from_pretrained(
 #     'jinaai/jina-vlm-v1',
-#     torch_dtype=torch.bfloat16,
+#     dtype=torch.bfloat16,
 #     attn_implementation='flash_attention_2',
 #     device_map='auto',
 #     trust_remote_code=True
@@ -317,7 +317,7 @@ conversation = [
                 'type': 'image',
                 'image': image,
             },
-            {'type': 'text', 'text': 'Describe this image.'},
+            {'type': 'text', 'text': 'Describe this image'},
         ],
     }
 ]
@@ -347,14 +347,10 @@ inputs = processor(
 
 # Move the inputs to the appropriate device and/or dtype
 device = torch.device('cuda')
-dtype = torch.float16
 model_inputs = {}
 for k, v in inputs.items():
     if isinstance(v, torch.Tensor):
-        if v.is_floating_point():
-            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
-        else:
-            model_inputs[k] = v.to(device, non_blocking=True)
+        model_inputs[k] = v.to(device, non_blocking=True)
     else:
         model_inputs[k] = v
 
@@ -362,7 +358,7 @@ for k, v in inputs.items():
 output = model.generate(
     **model_inputs,
     generation_config=GenerationConfig(
-        max_new_tokens=20, do_sample=False,
+        max_new_tokens=1024, do_sample=False,
     ),
     return_dict_in_generate=True,
     use_model_defaults=True,
@@ -390,7 +386,7 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     'jinaai/jina-vlm-v1',
     device_map='auto', 
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     attn_implementation='flash_attention_2',
     trust_remote_code=True
 )
@@ -441,7 +437,7 @@ for k, v in inputs.items():
 output = model.generate(
     **model_inputs,
     generation_config=GenerationConfig(
-        max_new_tokens=20, do_sample=False,
+        max_new_tokens=1024, do_sample=False,
     ),
     return_dict_in_generate=True,
     use_model_defaults=True,
@@ -468,7 +464,7 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     'jinaai/jina-vlm-v1',
     device_map='auto', 
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     attn_implementation='flash_attention_2',
     trust_remote_code=True
 )
@@ -508,7 +504,7 @@ for k, v in inputs.items():
 output = model.generate(
     **model_inputs,
     generation_config=GenerationConfig(
-        max_new_tokens=20, do_sample=False,
+        max_new_tokens=1024, do_sample=False,
     ),
     return_dict_in_generate=True,
     use_model_defaults=True,
@@ -535,7 +531,7 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     'jinaai/jina-vlm-v1',
     device_map='auto', 
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     attn_implementation='flash_attention_2',
     trust_remote_code=True
 )
@@ -599,7 +595,7 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModelForCausalLM.from_pretrained(
     'jinaai/jina-vlm-v1',
     device_map='auto', 
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     attn_implementation='flash_attention_2',
     trust_remote_code=True
 )
@@ -701,7 +697,7 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModel.from_pretrained(
     'jinaai/jina-vlm-v1',
     device_map='auto', 
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     attn_implementation='flash_attention_2',
     trust_remote_code=True
 )

blocks_jvlm.py

@@ -1294,6 +1294,7 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
         # image_features:
         # (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)
         bs, ncrops = image_features.shape[:2]
+        ogtype = image_features.dtype
 
         if self.padding_embed_type is not None:
             assert image_masks is not None
@@ -1322,6 +1323,7 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
                     partial_pad, -1
                 )
 
+        image_features = image_features.to(dtype=ogtype)
         image_features = self.feature_dropout(image_features)
         image_features = image_features.reshape((bs, ncrops) + self.n_patches + (-1,))
         pad_h = self.n_patches[0] % self.pooling_h

modeling_jvlm.py

@@ -388,6 +388,7 @@ class JinaVLMTextModel(JinaPreTrainedModel):
             batch_idx = torch.arange(bs, device=x.device)
             batch_idx = torch.tile(batch_idx[:, None], [1, image_features.shape[1]])
             image_features = image_features.to(x.device)
+            x = x.clone()  # Clone x to avoid in-place operation on leaf tensor
             x[batch_idx[valid], image_input_idx[valid]] += image_features[valid]
 
         if not self.rope:

test.py

@@ -0,0 +1,70 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+images = [
+    'https://picsum.photos/id/22/4434/3729',
+    'https://picsum.photos/id/49/1280/792'
+]
+conversations = [
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[0]},
+                {'type': 'text', 'text': 'What is the man doing in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[1]},
+                {'type': 'text', 'text': 'What country\'s flag is in this image?'},
+            ],
+        }
+    ],
+    
+]
+texts = processor.apply_chat_template(conversations, add_generation_prompt=True)
+inputs = processor(
+    text=texts,
+    images=images,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=1024, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
+)
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)