from transformers import ProcessorMixin, AutoImageProcessor, AutoTokenizer

class NexaVisionMoEProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
        if image_processor is None:
            image_processor = AutoImageProcessor.from_pretrained("google/siglip2-so400m-patch16-naflex")
        if tokenizer is None:
            tokenizer = AutoTokenizer.from_pretrained("Neura-Tech-AI/Nexa-AI-4x4B-Instruct")
        
        super().__init__(image_processor, tokenizer, **kwargs)

    def __call__(self, text=None, images=None, return_tensors="pt", **kwargs):
        if text is None and images is None:
            raise ValueError("You must provide either text or images.")

        output_kwargs = {}
        if images is not None:
            image_inputs = self.image_processor(images, return_tensors=return_tensors, **kwargs)
            output_kwargs.update(image_inputs)

        if text is not None:
            text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
            output_kwargs.update(text_inputs)

        return output_kwargs

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)