processing_longcat_next.py

import re
from typing import Union, List
from types import SimpleNamespace

import torch
import librosa
import soundfile as sf
import numpy as np
from transformers import AutoFeatureExtractor
from transformers.audio_utils import mel_filter_bank
from transformers.configuration_utils import PretrainedConfig
from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from transformers.processing_utils import (
    AudioKwargs,
    ImagesKwargs,
    ProcessingKwargs,
    ProcessorMixin,
    VideosKwargs,
)
from transformers.utils import logging

logger = logging.get_logger(__name__)


class LongcatNextProcessorKwargs(ProcessingKwargs, total=False):
    images_kwargs: ImagesKwargs
    videos_kwargs: VideosKwargs
    audio_kwargs: AudioKwargs
    _defaults = {
        "text_kwargs": {
            "padding": False,
            "padding_side": "left",
            "return_attention_mask": False,
        }
    }


class LongcatNextAudioProcessor(FeatureExtractionMixin):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mel_filters = mel_filter_bank(
            num_frequency_bins=1 + self.n_fft // 2,
            num_mel_filters=self.num_mel_bins,
            min_frequency=0.0,
            max_frequency=self.sampling_rate / 2.0,
            sampling_rate=self.sampling_rate,
            norm="slaney",
            mel_scale="slaney",
        )
        self.window = torch.hann_window(self.n_fft)

    @staticmethod
    def zero_mean_unit_var_norm(x):
        return (x - x.mean()) / torch.sqrt(x.var() + 1e-8)

    def load_audio_waveform(self, uri, metadata=None, waveform_tensor=None, return_tensors=True, do_normalize=False):
        if metadata is None or waveform_tensor is None:
            # 使用 librosa 统一处理所有音频格式（包括 mp3, wav, flac 等）
            # librosa.load 返回的已经是归一化的 float32 数据
            waveform_np, sample_rate = librosa.load(uri, sr=None, mono=False)

            # 转换为 tensor，确保维度为 (channels, samples)
            if waveform_np.ndim == 1:
                waveform_tensor = torch.from_numpy(waveform_np).unsqueeze(0)
            else:
                waveform_tensor = torch.from_numpy(waveform_np)

            # 获取音频元信息
            try:
                sf_info = sf.info(uri)
                metadata = SimpleNamespace(
                    sample_rate=sample_rate,
                    num_frames=waveform_tensor.shape[1],
                    num_channels=waveform_tensor.shape[0],
                    bits_per_sample=getattr(sf_info, 'bits_per_sample', 16),
                    encoding=getattr(sf_info, 'subtype', 'PCM_F')
                )
            except Exception:
                # 如果 soundfile.info 失败，使用 librosa 提供的信息
                metadata = SimpleNamespace(
                    sample_rate=sample_rate,
                    num_frames=waveform_tensor.shape[1],
                    num_channels=waveform_tensor.shape[0],
                    bits_per_sample=16,
                    encoding='PCM_F'
                )

        assert(metadata.num_channels <= 2), "acoustic file with {} channels.".format(metadata.num_channels)  # whisper only accept mono channel audio

        if self.sampling_rate != metadata.sample_rate:
            # 使用 torch.functional 进行重采样
            waveform_tensor = torch.nn.functional.interpolate(
                waveform_tensor.unsqueeze(0),
                size=int(waveform_tensor.shape[1] * self.sampling_rate / metadata.sample_rate),
                mode='linear',
                align_corners=False
            ).squeeze(0)

        # downmix to mono channel https://trac.ffmpeg.org/wiki/AudioChannelManipulation
        if metadata.num_channels > 1:
            waveform_tensor = torch.mean(waveform_tensor, dim=0, keepdim=True)

        # normalized to zero mean (Qwen Audio没有处理 但Whisper官方实现)
        if do_normalize:
            waveform_tensor = self.zero_mean_unit_var_norm(waveform_tensor)

        if return_tensors:  # (channels, samples)
            return waveform_tensor
        else:
            return waveform_tensor.numpy()

    def split_with_overlap(self, waveform):  # 如果长度超过最大长度限制 分割为带overlap的多段
        channels, wave_samples = waveform.shape
        max_audio_samples = self.max_audio_seconds * self.sampling_rate
        if wave_samples <= max_audio_samples or self.split_overlap < 0:
            return [waveform]  # 没有超出最大长度or截断逻辑 统一返回list

        split_waveform, start = [], 0
        while start < wave_samples:  # 统一按秒数对齐overlap
            if start > int(self.sampling_rate * self.split_overlap):
                start -= int(self.sampling_rate * self.split_overlap)  # 0表示没有overlap，>0 overlap对应秒数
            end = min(start + max_audio_samples, wave_samples)
            if end - start>= self.n_fft: # 保证至少有一帧数据
                split_waveform.append(waveform[:, start:end])  # 注意这里可能会切割出特别短的片段 需要在预处理判断并丢弃
            start = end
        return split_waveform

    @classmethod
    def inference_output_length(self, input_length, kernel_size, stride_size, avg_pooler):
        # for whisper + bridge
        encoder_length = (input_length + 2 * (kernel_size // 2) - kernel_size) // 1 + 1  # conv layer1 with pad=1
        encoder_length = (encoder_length + 2 * (kernel_size // 2) - kernel_size) // stride_size + 1  # conv layer2 with pad=1
        if avg_pooler > 1:
            bridge_length = encoder_length // avg_pooler
        return encoder_length, bridge_length

    def extract_fbank_features(self, waveform):
        # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/feature_extraction_whisper.py
        channels, wave_samples = waveform.shape
        assert(wave_samples >= self.n_fft)
        valid_frame_nums = min(self.max_audio_seconds * self.sampling_rate // self.hop_length, wave_samples // self.hop_length + 1)
        if wave_samples < self.max_audio_seconds * self.sampling_rate:
            waveform = torch.nn.functional.pad(waveform, (0, self.max_audio_seconds * self.sampling_rate - wave_samples), "constant", 0)
        else:
            waveform = waveform[:, :self.max_audio_seconds * self.sampling_rate]

        # window = torch.hann_window(self.n_fft)
        stft = torch.stft(waveform, self.n_fft, self.hop_length, window=self.window, return_complex=True)  # fft, len(wave) // n_fft // 2 + 1
        magnitudes = stft[..., :-1].abs() ** 2

        mel_filters = torch.from_numpy(self.mel_filters).type(torch.float32)
        mel_spec = mel_filters.T @ magnitudes
        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
        if waveform.dim() == 2:
            max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
            log_spec = torch.maximum(log_spec, max_val - 8.0)
        else:
            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
        log_spec = (log_spec + 4.0) / 4.0

        log_spec = log_spec[0].numpy()  # (channel, filters, samples) -> (filters, samples)
        log_spec[:, valid_frame_nums:] = 0.0  # pad0

        return log_spec, valid_frame_nums

    def process(self, audio_path, **kwargs):
        metadata, waveform_tensors = None, None
        waveforms = self.load_audio_waveform(audio_path, metadata, waveform_tensors, True)
        waveforms = self.split_with_overlap(waveforms)

        ret_audio, ret_encoder_length, ret_bridge_length = [], [], []
        for i, waveform in enumerate(waveforms):
            audio, input_length = self.extract_fbank_features(waveform)
            encoder_length, bridge_length = self.inference_output_length(input_length, self.kernel_size, self.stride_size, self.avg_pooler)
            if bridge_length <= 0:
                continue

            ret_audio.append(audio)
            ret_encoder_length.append(encoder_length)
            ret_bridge_length.append(bridge_length)
        return ret_audio, ret_encoder_length, ret_bridge_length

    def __call__(self, audio: Union[str, List[str]], **kwargs):
        if isinstance(audio, str):
            audio = [audio]
        results = {
            "audio": [],
            "encoder_length": [],
            "bridge_length": [],
        }
        for audio_path in audio:
            audio, encoder_length, bridge_length = self.process(audio_path, **kwargs)
            results["audio"].append(audio)
            results["encoder_length"].append(encoder_length)
            results["bridge_length"].append(bridge_length)
        return results


class LongcatNextProcessor(ProcessorMixin):

    attributes = ["image_processor", "video_processor", "audio_processor", "tokenizer"]

    image_processor_class = "Qwen2VLImageProcessor"
    video_processor_class = "Qwen2VLImageProcessor"
    audio_processor_class = "LongcatNextAudioProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, video_processor=None, audio_processor=None, tokenizer=None, chat_template=None, **kwargs):
        super().__init__(image_processor, video_processor, audio_processor, tokenizer, chat_template=chat_template)
        init_token_list = [
            "image_start_token", "image_end_token", "image_pad_token", "image_newline_token",
            "audio_start_token", "audio_end_token", "audio_pad_token",
        ]
        for attr in init_token_list:
            token_str = self.tokenizer.init_kwargs.get(attr)
            token_ids = self.tokenizer.encode(token_str, add_special_tokens=False)
            assert len(token_ids) == 1, (f"{attr}='{token_str}' encode to get {len(token_ids)} id(s) {token_ids}, expect 1 id")
            setattr(self, f"{attr}", token_str)
            setattr(self, f"{attr}_id", token_ids[0])

    def __call__(
        self,
        text: str,
        **kwargs,
    ) -> List["LongcatNextProcessorOutput"]:

        if text is None:
            raise ValueError("You need to specify either a `text` input to process.")

        output_kwargs = self._merge_kwargs(
            LongcatNextProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        assert isinstance(text, str)

        image_path_list = re.findall(rf"{self.image_start_token}(.*?){self.image_end_token}", text)
        audio_path_list = re.findall(rf"{self.audio_start_token}(.*?){self.audio_end_token}", text)

        if len(image_path_list) > 0:
            images_inputs = self.image_processor(images=image_path_list, **output_kwargs["images_kwargs"])
            image_grid_thw = images_inputs["image_grid_thw"]
            for i, image_path in enumerate(image_path_list):
                image_token_num = image_grid_thw[i][0] * (image_grid_thw[i][1]//self.image_processor.spatial_merge_size) * (image_grid_thw[i][2]//self.image_processor.spatial_merge_size)
                text = text.replace(f"{self.image_start_token}{image_path}{self.image_end_token}", f"{self.image_start_token}{self.image_pad_token * image_token_num}{self.image_end_token}")
        else:
            images_inputs = {}

        if len(audio_path_list) > 0:
            audio_inputs = self.audio_processor(audio=audio_path_list, **output_kwargs["audio_kwargs"])
            for i, audio_path in enumerate(audio_path_list):
                audio_token_num = np.sum(audio_inputs["bridge_length"][i])
                text = text.replace(f"{self.audio_start_token}{audio_path}{self.audio_end_token}", f"{self.audio_start_token}{self.audio_pad_token * audio_token_num}{self.audio_end_token}")
            for key in audio_inputs:
                audio_inputs[key] = [val for b_val in audio_inputs[key] for val in b_val]
        else:
            audio_inputs = {}

        texts_inputs = self.tokenizer([text], **output_kwargs["text_kwargs"])

        batch_feature_func = lambda x: BatchFeature(
            data={**x},
            tensor_type=kwargs.get("return_tensors"),
        )
        return (
            batch_feature_func(texts_inputs),
            batch_feature_func({k.replace("image", "visual"): v for k, v in images_inputs.items()}) if len(images_inputs) > 0 else None,
            batch_feature_func(audio_inputs) if len(audio_inputs) > 0 else None,
        )


class LongcatNextAudioProcessorConfig(PretrainedConfig):
    pass
AutoFeatureExtractor.register(LongcatNextAudioProcessorConfig, LongcatNextAudioProcessor)


__all__ = ["LongcatNextAudioProcessor", "LongcatNextProcessor"]