""" This module defines a self-contained StarVector model with support for remote code loading. """ import os import torch import torch.nn as nn from transformers import PreTrainedModel, PretrainedConfig from typing import Optional, Union, List from abc import ABC, abstractmethod # Import components - these will be included in the HF repo from .starvector.image_encoder import ImageEncoder from .starvector.adapter import Adapter # === Model Configuration === class StarVectorConfig(PretrainedConfig): model_type = "starvector" def __init__( self, starcoder_model_name: str = "bigcode/starcoderbase-1b", image_encoder_type: str = "clip", adapter_norm: str = "layer_norm", image_size: int = 224, max_length: int = 8192, max_length_train: int = 8192, use_flash_attn: bool = True, use_cache: bool = True, num_attention_heads: int = 16, num_hidden_layers: int = 24, vocab_size: int = 49152, hidden_size: int = 2048, num_kv_heads: int = 4, torch_dtype: str = "bfloat16", **kwargs, ): # Initialize the parent config first super().__init__(**kwargs) self.starcoder_model_name = starcoder_model_name self.image_encoder_type = image_encoder_type self.adapter_norm = adapter_norm self.image_size = image_size self.max_length = max_length self.max_length_train = max_length_train self.use_flash_attn = use_flash_attn self.use_cache = use_cache self.num_attention_heads = num_attention_heads self.num_hidden_layers = num_hidden_layers self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_kv_heads = num_kv_heads self.torch_dtype = torch_dtype # === Base Model Classes === class StarVectorBase(nn.Module, ABC): def __init__(self, config, **kwargs): super().__init__() self.task = kwargs.get('task', 'im2svg') self.model_precision = kwargs.get('model_precision', config.torch_dtype) # Instantiate the SVG transformer using the abstract method. self.svg_transformer = self._get_svg_transformer(config, **kwargs) if self.use_image_encoder(): self.image_encoder = ImageEncoder(config, **kwargs) self.image_projection = self.get_adapter(config, **kwargs).to(dtype=self.model_precision) else: self.query_length = 0 self.max_length = config.max_length_train - getattr(self, "query_length", 0) - 4 self.train_image_encoder = kwargs.get('train_image_encoder', False) self.train_LLM = kwargs.get('train_LLM', False) self._freeze_parameters(self.train_image_encoder, self.train_LLM) @abstractmethod def _get_svg_transformer(self, config, **kwargs): """Get SVG transformer model - implementation differs between versions""" pass def _freeze_parameters(self, train_image_encoder, train_LLM): if self.use_image_encoder(): for _, param in self.image_encoder.named_parameters(): param.requires_grad = train_image_encoder for _, param in self.image_projection.named_parameters(): param.requires_grad = train_image_encoder for _, param in self.svg_transformer.transformer.named_parameters(): param.requires_grad = train_LLM def use_image_encoder(self): return self.task == 'im2svg' def get_adapter(self, config, **kwargs): # Determine hidden size and query length based on the image encoder type. if config.image_encoder_type == 'clip': hidden_size = self.image_encoder.num_features self.query_length = 257 elif config.image_encoder_type == 'vqgan': hidden_size = 256 self.query_length = 196 else: hidden_size = 256 # default fallback self.query_length = 200 llm_hidden_size = config.hidden_size # assuming the transformer hidden size return Adapter(hidden_size, llm_hidden_size, adapter_norm=config.adapter_norm, query_length=self.query_length, dropout_prob=kwargs.get('dropout', 0.1)) def forward(self, batch): # Simplified forward pass where we assume batch has an "image" key. image = batch["image"] if self.use_image_encoder(): embedded_image = self.image_encoder(image) conditioning_embeds = self.image_projection(embedded_image) # For demo purposes, we generate dummy input embeddings (replace with your logic) inputs_embeds = self.svg_transformer.transformer.wte( torch.randint(0, self.svg_transformer.transformer.wte.num_embeddings, (image.size(0), self.max_length)) ) else: inputs_embeds = self.svg_transformer.transformer.wte( torch.randint(0, self.svg_transformer.transformer.wte.num_embeddings, (image.size(0), self.max_length)) ) return inputs_embeds # Dummy return def generate_im2svg(self, batch, **kwargs): # Prepare generation inputs (dummy implementation) image = batch["image"] if self.use_image_encoder(): embedded_image = self.image_encoder(image) conditioning_embeds = self.image_projection(embedded_image) else: conditioning_embeds = torch.zeros((image.size(0), 10, 1), device=image.device) generation_output = self.svg_transformer.transformer.generate(inputs_embeds=conditioning_embeds, max_length=kwargs.get('max_length', 30)) raw_svg = self.svg_transformer.tokenizer.batch_decode(generation_output, skip_special_tokens=True) return raw_svg @abstractmethod def _get_embeddings(self, input_ids): """Get embeddings from input ids - implementation differs between v1 and v2""" pass @abstractmethod def _get_svg_text(self, svg_list): """Get SVG text with appropriate end tokens - implementation differs between v1 and v2""" pass # V1 implementation: Delegates transformer creation to the external LLM file. class StarVectorStarCoder(StarVectorBase): def __init__(self, config, **kwargs): super().__init__(config, **kwargs) def _get_svg_transformer(self, config, **kwargs): from starvector.model.llm.starcoder import StarCoderModel # V1: use StarCoderModel from external file return StarCoderModel(config, **kwargs) def _get_embeddings(self, input_ids): """V1-specific embedding method""" # This follows the implementation in starvector/model/models/starvector_v1.py. return self.svg_transformer.transformer.transformer.wte(input_ids) def _get_svg_text(self, svg_list): """V1-specific SVG text preparation""" return [t + self.svg_transformer.tokenizer.eos_token for t in svg_list] # V2 implementation: Delegates transformer creation to the external V2 LLM file. class StarVectorStarCoder2(StarVectorBase): def __init__(self, config, **kwargs): super().__init__(config, **kwargs) def _get_svg_transformer(self, config, **kwargs): from starvector.model.llm.starcoder2 import StarCoderModel # V2: use external StarCoderModel from starcoder2.py return StarCoderModel(config, **kwargs) def _get_embeddings(self, input_ids): """V2-specific embedding method""" return self.svg_transformer.transformer.model.embed_tokens(input_ids) def _get_svg_text(self, svg_list): """V2-specific SVG text preparation""" return [t + self.svg_transformer.svg_end_token + self.svg_transformer.tokenizer.eos_token for t in svg_list] def _get_im2svg_specific_kwargs(self, kwargs): """V2-specific generation kwargs""" return { 'eos_token_id': self.svg_transformer.svg_end_token_id, } def _get_text2svg_specific_kwargs(self, kwargs): """V2-specific text2svg generation kwargs""" return { 'eos_token_id': self.svg_transformer.tokenizer.eos_token_id, } # === Main Model Class for Hugging Face === class StarVectorForCausalLM(PreTrainedModel): config_class = StarVectorConfig _no_split_modules = [] def __init__(self, config, **kwargs): super().__init__(config) # Choose V2 if the model name indicates starcoder2; otherwise use V1. if "starcoder2" in config.starcoder_model_name.lower(): self.model = StarVectorStarCoder2(config=config, **kwargs) else: self.model = StarVectorStarCoder(config=config, **kwargs) def forward(self, batch): return self.model(batch) def generate_im2svg(self, batch, **kwargs): return self.model.generate_im2svg(batch, **kwargs) def generate_im2text(self, batch, **kwargs): return self.model.generate_im2text(batch, **kwargs) def process_images(self, images): return self.model.image_encoder.process_images(images) # === Registration for Autonomous Loading === StarVectorConfig.register_for_auto_class() StarVectorForCausalLM.register_for_auto_class("AutoModelForCausalLM")