Update modeling_glm.py

modeling_glm.py

@@ -417,7 +417,7 @@ class GlmSdpaAttention(GlmAttention):
             )
 
         bsz, q_len, _ = hidden_states.size()
-
+        
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
@@ -425,7 +425,7 @@ class GlmSdpaAttention(GlmAttention):
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
+        
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(
             query_states, key_states, cos, sin, partial_rotary_factor=self.partial_rotary_factor
@@ -763,21 +763,36 @@ class GlmModel(GlmPreTrainedModel):
             assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
             inputs_embeds = self.embed_tokens(input_ids)
             new_input_embeds = []
-            boi_token_flags = [True if self.config.boi_token_id in input_id.tolist() else False for input_id in input_ids]
-            if is_empty(images):
-                images = torch.zeros([1, 3, 672, 672]).to(input_ids.device)
-            images_features = self.vision(images).to(device=inputs_embeds.device, dtype=inputs_embeds.dtype)
+            multi_flags = [True if self.config.boi_token_id in input_id.tolist() else False for input_id in input_ids]
+            images_features = None
+            if not is_empty(images) and images.bool().any():
+                imgs = list()
+                for i in range(len(multi_flags)):
+                    if multi_flags[i]:
+                        imgs.append(images[i])
+                imgs = torch.stack(imgs, dim=0)
+            else:
+                imgs = torch.unsqueeze(images[0], 0)
+            images_features = self.vision(imgs).to(inputs_embeds.dtype)
             image_count = 0
             for i in range(len(input_ids)):
                 input_id = input_ids[i].tolist()
-                if boi_token_flags[i]:
+                if multi_flags[i]:
                     boi_token_pos = input_id.index(self.config.boi_token_id)
                     assert boi_token_pos >= 0, "begin_of_image not found!"
                     num_image_padding_tokens = input_id.count(self.config.boi_token_id)
-                    assert num_image_padding_tokens == images_features[image_count].shape[0], f"Wrong image padding token number: {num_image_padding_tokens}"
-                    new_input_embeds.append(torch.cat(
-                        (inputs_embeds[i, :boi_token_pos], images_features[image_count],
-                        inputs_embeds[i, boi_token_pos + num_image_padding_tokens:])))
+                    assert (
+                        num_image_padding_tokens == images_features[image_count].shape[0]
+                    ), f"Wrong image padding token number: {num_image_padding_tokens}"
+                    new_input_embeds.append(
+                        torch.cat(
+                            (
+                                inputs_embeds[i, :boi_token_pos],
+                                images_features[image_count].to(inputs_embeds.device),
+                                inputs_embeds[i, boi_token_pos + num_image_padding_tokens :],
+                            )
+                        )
+                    )
                     image_count += 1
                 else:
                     new_input_embeds.append(inputs_embeds[i] + (0 * images_features[0].sum()))
@@ -1316,4 +1331,4 @@ __all__ = [
     "GlmModel",
     "GlmForCausalLM",
     "GlmForSequenceClassification",
-]
+]