Update starvector_arch.py

starvector_arch.py

@@ -35,7 +35,6 @@ class SimpleStarVectorProcessor(ProcessorMixin):
         self.mean = mean
         self.std = std
         self.size = size
-        
         self.normalize = transforms.Normalize(mean=mean, std=std)        
         
         self.transform = transforms.Compose([
@@ -50,7 +49,7 @@ class SimpleStarVectorProcessor(ProcessorMixin):
         super().__init__(tokenizer=tokenizer)
 
 
-    def __call__(self, images=None, text=None, **kwargs) -> BatchFeature:
+    def __call__(self, images=None, text=None, max_length=None, **kwargs) -> BatchFeature:
         """
         Process images and/or text inputs.
         
@@ -65,16 +64,32 @@ class SimpleStarVectorProcessor(ProcessorMixin):
         image_inputs = {}
         if images is not None:
             if isinstance(images, (list, tuple)):
-                images_ = [self.transform(img) for img in images]
+                images_ = torch.stack([self.transform(img) for img in images])
             else:
                 images_ = self.transform(images)
             image_inputs = {"pixel_values": images_}
         
         text_inputs = {}
         if text is not None:
-            text_inputs = self.tokenizer(text, **kwargs)
+            text_inputs = self.tokenizer(
+                text, truncation=True, 
+                add_special_tokens=True, 
+                padding='longest', 
+                max_length=max_length, 
+                return_tensors="pt"
+            )
+
         return BatchFeature(data={**text_inputs, **image_inputs})
 
+    def _pad_to_square(self, img):
+        # Calculate padding to make the image square
+        width, height = img.size
+        max_dim = max(width, height)
+        padding = [(max_dim - width) // 2, (max_dim - height) // 2]
+        padding += [max_dim - width - padding[0], max_dim - height - padding[1]]
+        return pad(img, padding, fill=255)  # Assuming white padding
+
+
 AutoProcessor.register(SimpleStarVectorProcessor, SimpleStarVectorProcessor)
 
 
@@ -128,6 +143,7 @@ class StarVectorForCausalLM(PreTrainedModel):
         else:
             from starvector.model.models.starvector_v1 import StarVectorStarCoder
             self.model = StarVectorStarCoder(config=config, **kwargs)
+            
 
     @property
     def supports_gradient_checkpointing(self):
@@ -142,70 +158,28 @@ class StarVectorForCausalLM(PreTrainedModel):
         if hasattr(self.model, 'svg_transformer') and hasattr(self.model.svg_transformer, 'gradient_checkpointing_enable'):
             self.model.svg_transformer.gradient_checkpointing_enable()
 
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        num_logits_to_keep: int = 0,
-    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
+    def forward(self,  inputs_embeds, input_ids, num_generations, num_logits_to_keep) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
         r"""
-        labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
-            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
-            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        Wrapper for the forward pass of the model.
         """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        device = inputs_embeds.device
 
-        transformer_outputs = self.model.svg_transformer.transformer(
-            input_ids,
-            past_key_values=past_key_values,
+        completion_embeds = self.model._get_embeddings(input_ids)
+        inputs_embeds = torch.cat([inputs_embeds.repeat(num_generations, 1, 1), completion_embeds], dim=1)
+        attention_mask = torch.ones_like(inputs_embeds[:, :, 0]).to(device)
+
+        transformer_outputs = self.model.svg_transformer.transformer.transformer(
+            inputs_embeds=inputs_embeds, 
             attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_attention_mask=encoder_attention_mask,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
         )
         hidden_states = transformer_outputs[0]
 
         # If GRPO requested only the last tokens, slice accordingly.
         if num_logits_to_keep > 0:
-            lm_logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+            lm_logits = self.model.svg_transformer.transformer.lm_head(hidden_states[:, -num_logits_to_keep:, :])
         else:
-            lm_logits = self.lm_head(hidden_states)
-
-        # lm_logits = self.lm_head(hidden_states)
-
+            lm_logits = self.model.svg_transformer.transformer.lm_head(hidden_states)
         loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
-        if not return_dict:
-            output = (lm_logits,) + transformer_outputs[1:]
-            return ((loss,) + output) if loss is not None else output
-
         return CausalLMOutputWithCrossAttentions(
             loss=loss,
             logits=lm_logits,
@@ -214,9 +188,6 @@ class StarVectorForCausalLM(PreTrainedModel):
             attentions=transformer_outputs.attentions,
             cross_attentions=transformer_outputs.cross_attentions,
         )
-    
-    # def forward(self, batch):
-    #     return self.model(batch)
 
     def generate_im2svg(self, batch, **kwargs):
         return self.model.generate_im2svg(batch, **kwargs)