Add HuggingFace format support for OmniLottie

decoder_hf.py

@@ -0,0 +1,179 @@
+import torch
+import torch.nn as nn
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig, PreTrainedModel
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLCausalLMOutputWithPast
+from typing import Any, Dict, List, Optional, Tuple, Union
+import os
+
+import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as qwen_modeling
+
+from configuration_lottie_decoder import LottieDecoderConfig
+
+
+class LottieDecoder(PreTrainedModel):
+    """
+    Autoregressive generative model for OmniLottie
+
+    Lottie animation generation model based on Qwen2.5-VL,
+    supports generating Lottie JSON code from videos.
+    """
+
+    config_class = LottieDecoderConfig
+    base_model_prefix = "lottie_decoder"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: LottieDecoderConfig):
+        """
+        Initialize LottieDecoder model
+
+        Args:
+            config (LottieDecoderConfig): Model configuration object
+        """
+        super().__init__(config)
+
+        self.config = config
+        self.pix_len = config.pix_len
+        self.text_len = config.text_len
+        self.vocab_size = config.vocab_size
+        self.bos_token_id = config.bos_token_id
+        self.eos_token_id = config.eos_token_id
+        self.pad_token_id = config.pad_token_id
+
+        print(f"Initializing LottieDecoder with base model: {config.base_model_path}")
+
+        # Create base model configuration
+        qwen_config = AutoConfig.from_pretrained(
+            config.base_model_path,
+            vocab_size=self.vocab_size,
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id,
+            trust_remote_code=True
+        )
+
+        # Load base Qwen2.5-VL model
+        self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            config.base_model_path,
+            config=qwen_config,
+            torch_dtype=getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype,
+            attn_implementation=config.attn_implementation,
+            ignore_mismatched_sizes=True
+        )
+
+        # Extend vocabulary to support Lottie tokens
+        self.transformer.resize_token_embeddings(self.vocab_size)
+
+        # Set to training mode initially (same as original decoder)
+        self.train()
+
+        print(f"LottieDecoder initialized successfully. Vocab size: {self.vocab_size}")
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """
+        Load LottieDecoder from pretrained model path
+
+        Supports two loading methods:
+        1. Load from Hugging Face standard format (recommended)
+        2. Load from old format pytorch_model.bin (backward compatible)
+        """
+        # Check if it's old format (contains pytorch_model.bin)
+        if os.path.isdir(pretrained_model_name_or_path):
+            old_format_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
+            if os.path.exists(old_format_path) and not os.path.exists(os.path.join(pretrained_model_name_or_path, 'config.json')):
+                print(f"Detected old format model, loading from {old_format_path}...")
+                return cls._from_old_format(pretrained_model_name_or_path, **kwargs)
+
+        # Use standard Hugging Face loading process
+        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+    @classmethod
+    def _from_old_format(cls, checkpoint_path, **kwargs):
+        """
+        Load model from old format (pytorch_model.bin)
+
+        Args:
+            checkpoint_path: Directory path containing pytorch_model.bin
+        """
+        # Extract configuration parameters
+        pix_len = kwargs.pop('pix_len', 4560)
+        text_len = kwargs.pop('text_len', 1500)
+        base_model_path = kwargs.pop('base_model_path', 'Qwen/Qwen2.5-VL-3B-Instruct')
+
+        # Create configuration
+        config = LottieDecoderConfig(
+            pix_len=pix_len,
+            text_len=text_len,
+            base_model_path=base_model_path
+        )
+
+        # Initialize model
+        model = cls(config)
+
+        # Load weights
+        model_file = os.path.join(checkpoint_path, 'pytorch_model.bin')
+        if os.path.exists(model_file):
+            state_dict = torch.load(model_file, map_location='cpu')
+            model.load_state_dict(state_dict, strict=False)
+            print(f"Successfully loaded weights from {model_file}")
+        else:
+            print(f"Warning: Model file not found {model_file}")
+
+        return model
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        pixel_values=None,
+        image_grid_thw=None,
+        pixel_values_videos=None,
+        video_grid_thw=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=False,
+        **kwargs
+    ):
+        """
+        Forward pass - currently for inference only, needs implementation for training
+        """
+        return self.transformer(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            labels=labels,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            **kwargs
+        )
+
+    def generate(self, *args, **kwargs):
+        """
+        Generate Lottie tokens
+
+        Directly calls the underlying transformer's generate method
+        """
+        return self.transformer.generate(*args, **kwargs)
+
+    def get_input_embeddings(self):
+        """Get input embeddings"""
+        return self.transformer.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        """Set input embeddings"""
+        self.transformer.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        """Get output embeddings"""
+        return self.transformer.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        """Set output embeddings"""
+        self.transformer.set_output_embeddings(new_embeddings)
+
+    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None):
+        """Resize token embeddings"""
+        return self.transformer.resize_token_embeddings(new_num_tokens)