SheetSage2/modeling_sheetsage2.py

"""Hugging Face SheetSage2 model with a shared MERT-v2 encoder."""

from dataclasses import dataclass
import copy
import hashlib
from pathlib import Path
import shutil
from types import SimpleNamespace
from typing import Optional, Tuple

import torch
from torch import nn
from torch.nn import functional as F
from transformers import AutoModel, BartConfig, PreTrainedModel
from transformers.models.bart.modeling_bart import BartDecoder
from transformers.utils import ModelOutput
from transformers.utils.hub import cached_file

from .configuration_mert2 import MERT2Config
from .modeling_mert2 import MERT2Model
from .configuration_sheetsage2 import SheetSage2Config
from .tokenization_sheetsage2 import SheetSage2Tokenizer


PROJECTIONS = ("query_proj", "key_proj", "value_proj", "out_proj")
BASE_CODE_HASHES = {
    "configuration_mert2.py": "77b53ec9d7ee31a599d744fb006e812c7eeaf7390deb46e2f460cf8c17b00bd6",
    "modeling_mert2.py": "b1a3174e5649c4b26b0c90d8626f0adacfbbba111a58ed3bb72ad651945a2f5c",
}


def _hf_relative_dependencies():
    # Transformers 4.45 copies direct imports into a fresh local module cache.
    from .audio_sheetsage2 import load_audio
    from .durations_sheetsage2 import DURATION_TEMPLATES
    from .exports_sheetsage2 import export_result
    from .io_sheetsage2 import atomic_write_text
    from .labels_sheetsage2 import STRUCTURE_LABELS
    from .midi_sheetsage2 import export_playback
    from .notation_sheetsage2 import generate_abc_from_exports
    from .rendering_sheetsage2 import render_outputs
    from .schema_sheetsage2 import get_prompt_multitask_schema
    from .tensors_sheetsage2 import WindowTensorWriter


def _sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(8 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


@dataclass
class SheetSage2EncoderOutput(ModelOutput):
    """Frame features. Block states exclude the separately returned input state."""

    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    backbone_last_hidden_state: Optional[torch.FloatTensor] = None
    mixed_hidden_state: Optional[torch.FloatTensor] = None
    input_hidden_state: Optional[torch.FloatTensor] = None
    backbone_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    feature_attention_mask: Optional[torch.BoolTensor] = None

    @property
    def last_hidden_state(self):
        return self.encoder_last_hidden_state


@dataclass
class SheetSage2Output(ModelOutput):
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Tuple] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    backbone_last_hidden_state: Optional[torch.FloatTensor] = None
    mixed_hidden_state: Optional[torch.FloatTensor] = None
    input_hidden_state: Optional[torch.FloatTensor] = None
    backbone_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    feature_attention_mask: Optional[torch.BoolTensor] = None


class AttentionAdapter(nn.Module):
    def __init__(self, width, rank):
        super().__init__()
        self.lora_A = nn.Linear(width, rank, bias=False)
        self.lora_B = nn.Linear(rank, width, bias=False)


class EncoderAdapters(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layers = nn.ModuleList()
        for _ in range(config.backbone_config["num_hidden_layers"]):
            layer = nn.Module()
            layer.attn = nn.Module()
            for name in PROJECTIONS:
                setattr(layer.attn, name, AttentionAdapter(config.backbone_config["hidden_size"], config.lora_rank))
            self.layers.append(layer)


class SheetSage2Model(PreTrainedModel):
    """Audio-to-symbolic model, with optional frame and decoder representations.

    ``from_pretrained`` loads the pinned MERT-v2 parent and merges attention
    adapters in float32. ``save_pretrained`` writes an independent merged model.
    ``forward`` returns raw vocabulary logits; grammar-masked generation scores
    are available separately through ``generate(output_scores=True)``.
    """

    config_class = SheetSage2Config
    base_model_prefix = ""
    main_input_name = "input_values"
    _supports_sdpa = True
    _tied_weights_keys = ["decoder.embed_tokens.weight", "output_projection.weight"]
    _no_split_modules = ["ConformerBlock", "BartDecoderLayer"]

    def __init__(self, config):
        super().__init__(config)
        self.hparams = SimpleNamespace(input_audio_length=config.input_audio_length, time_hz=config.time_hz)
        self.max_output_seq_len = config.max_output_seq_len
        self.tokenizer = SheetSage2Tokenizer(
            config.input_audio_length, config.time_hz, config.tokenizer_schema_version,
            expected_fingerprint=config.tokenizer_fingerprint,
        )
        if self.tokenizer.n_tokens != config.vocab_size:
            raise ValueError("Tokenizer vocabulary size does not match the model.")
        if config.weights_format == "adapter":
            self.encoder = None
            self.adapter = EncoderAdapters(config)
        else:
            ec = MERT2Config(**config.backbone_config)
            ec._attn_implementation = config.encoder_attn_implementation
            self.encoder = MERT2Model(ec)
            self.adapter = None
        self.layer_weight = nn.Parameter(torch.zeros(config.backbone_config["num_hidden_layers"] + 1))
        self.encoder_projection = nn.Linear(config.backbone_config["hidden_size"], config.hidden_size)
        dc = BartConfig(
            vocab_size=config.vocab_size, d_model=config.hidden_size,
            decoder_layers=config.decoder_layers, decoder_attention_heads=config.num_attention_heads,
            decoder_ffn_dim=config.intermediate_size, max_position_embeddings=config.max_output_seq_len,
            dropout=config.decoder_dropout, attention_dropout=config.decoder_dropout,
            activation_dropout=config.decoder_dropout, activation_function="gelu",
            pad_token_id=config.pad_token_id, bos_token_id=config.bos_token_id,
            eos_token_id=config.eos_token_id, is_encoder_decoder=True, use_cache=True,
        )
        dc._attn_implementation = "sdpa"
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.decoder = BartDecoder(dc, embed_tokens=self.token_embedding)
        self.decoder.gradient_checkpointing_disable()
        self.output_projection = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.output_projection.weight = self.token_embedding.weight
        self.lora_merged = config.weights_format == "merged"
        self.post_init()

    def get_input_embeddings(self):
        return self.token_embedding

    def set_input_embeddings(self, value):
        self.token_embedding = value
        self.decoder.embed_tokens.weight = value.weight

    def tie_weights(self):
        super().tie_weights()
        if hasattr(self, "decoder") and hasattr(self, "token_embedding"):
            self.decoder.embed_tokens.weight = self.token_embedding.weight

    def get_output_embeddings(self):
        return self.output_projection

    def set_output_embeddings(self, value):
        self.output_projection = value

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)
            if isinstance(module, nn.Embedding) and module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    @torch.no_grad()
    def merge_lora(self):
        if self.lora_merged:
            return self
        if self.encoder is None or self.adapter is None:
            raise RuntimeError("Load both the MERT-v2 parent and adapters before merging.")
        scale = self.config.lora_alpha / self.config.lora_rank
        with torch.autocast("cpu", enabled=False):
            for layer, adapter in zip(self.encoder.layers, self.adapter.layers):
                for name in PROJECTIONS:
                    projection = getattr(layer.attn, name)
                    update = getattr(adapter.attn, name)
                    values = (projection.weight, update.lora_A.weight, update.lora_B.weight)
                    if any(value.device.type != "cpu" or value.dtype != torch.float32 for value in values):
                        raise ValueError("Merge adapters on CPU with float32 parameters.")
                    projection.weight.add_((update.lora_B.weight @ update.lora_A.weight) * scale)
        self.adapter = None
        self.lora_merged = True
        self.config.weights_format = "merged"
        return self

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        base_path = kwargs.pop("base_model_path", None)
        requested_dtype = kwargs.pop("torch_dtype", torch.float32)
        target_device = kwargs.pop("device_map", None)
        backend = kwargs.pop("attn_implementation", None)
        return_loading_info = kwargs.pop("output_loading_info", False)
        config = kwargs.pop("config", None)
        hub_keys = ("cache_dir", "force_download", "local_files_only", "token", "revision", "subfolder")
        hub_args = {name: kwargs[name] for name in hub_keys if name in kwargs}
        if config is None:
            config = cls.config_class.from_pretrained(pretrained_model_name_or_path, **hub_args)
        if backend is not None:
            config.encoder_attn_implementation = backend
        if config.encoder_attn_implementation not in {"sdpa", "flash_attention_2"}:
            raise ValueError("Select attn_implementation='sdpa' or 'flash_attention_2'.")
        if isinstance(target_device, dict):
            if set(target_device) != {""}:
                raise ValueError("Use a single device for SheetSage2: device_map={'': 'cuda:0'}.")
            target_device = target_device[""]
        if target_device == "auto":
            target_device = "cuda" if torch.cuda.is_available() else "cpu"
        if requested_dtype == "auto":
            requested_dtype = torch.float32
        if isinstance(requested_dtype, str):
            requested_dtype = getattr(torch, requested_dtype, None)
        if requested_dtype not in {None, torch.float32, torch.bfloat16, torch.float16}:
            raise ValueError("Use torch_dtype float32, bfloat16, float16, or 'auto'.")
        # Adapters must be loaded and merged before any reduced-precision cast.
        model, loading_info = super().from_pretrained(
            pretrained_model_name_or_path, *model_args, config=config,
            torch_dtype=torch.float32, attn_implementation="sdpa", output_loading_info=True, **kwargs,
        )
        if any(loading_info.get(name) for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
            raise ValueError(f"Incomplete or incompatible SheetSage2 weights: {loading_info}")
        if config.weights_format == "adapter":
            parent = str(base_path or config.base_model_name_or_path)
            parent_hub_args = {name: hub_args[name] for name in ("cache_dir", "force_download", "local_files_only", "token") if name in hub_args}
            parent_hub_args["revision"] = config.base_model_revision
            expected_files = dict(BASE_CODE_HASHES, **{"model.safetensors": config.base_model_sha256})
            for filename, expected in expected_files.items():
                resolved = cached_file(parent, filename, **parent_hub_args)
                if _sha256(resolved) != expected:
                    raise ValueError(f"MERT-v2 parent integrity check failed: {filename}")
            model.encoder = AutoModel.from_pretrained(
                parent, trust_remote_code=True, code_revision=config.base_model_revision,
                torch_dtype=torch.float32, attn_implementation=config.encoder_attn_implementation,
                **parent_hub_args,
            )
            for name, expected in config.backbone_config.items():
                if name in ("hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads",
                            "sampling_rate", "hop_length", "n_fft", "win_length", "num_mel_bins",
                            "conv_depthwise_kernel_size", "rotary_embedding_base", "subsampling_channels",
                            "subsampling_depths", "layer_norm_eps", "subsampling_layer_norm_eps", "variant"):
                    if getattr(model.encoder.config, name) != expected:
                        raise ValueError(f"MERT-v2 parent architecture mismatch: {name}")
            model.merge_lora()
        if requested_dtype is not None:
            model.to(dtype=requested_dtype)
        if target_device is not None:
            model.to(target_device)
        # Resource lookup follows the caller's cache and offline settings. These
        # transient options are never written into config or model weights.
        model._hub_resource_options = {name: hub_args[name] for name in ("cache_dir", "local_files_only", "token") if name in hub_args}
        resource_args = dict(model._hub_resource_options, local_files_only=True,
                             revision=config._commit_hash or hub_args.get("revision"),
                             subfolder=hub_args.get("subfolder", ""),
                             _raise_exceptions_for_missing_entries=False)
        resource_config = cached_file(str(pretrained_model_name_or_path), "config.json", **resource_args)
        model._source_snapshot = Path(resource_config).parent if resource_config else Path(pretrained_model_name_or_path)
        model.eval().requires_grad_(False)
        return (model, loading_info) if return_loading_info else model

    def save_pretrained(self, save_directory, *args, **kwargs):
        if not self.lora_merged or self.encoder is None:
            raise ValueError("Load and merge the model before saving a standalone snapshot.")
        original_config = self.config
        source = original_config._name_or_path
        self.config = copy.deepcopy(original_config)
        self.config.weights_format = "merged"
        self.config._name_or_path = ""
        self.config.backbone_config.pop("_name_or_path", None)
        try:
            result = super().save_pretrained(save_directory, *args, **kwargs)
            from .processing_sheetsage2 import SheetSage2Processor
            SheetSage2Processor.from_model_config(self.config).save_pretrained(save_directory)
        finally:
            self.config = original_config
        source_dir = getattr(self, "_source_snapshot", Path(source))
        if source and not source_dir.is_dir():
            from huggingface_hub import snapshot_download
            try:
                source_dir = Path(snapshot_download(source, revision=original_config._commit_hash,
                                                     cache_dir=getattr(self, "_hub_resource_options", {}).get("cache_dir"),
                                                     local_files_only=True))
            except (OSError, ValueError):
                source_dir = None
        if source and source_dir is not None and source_dir.is_dir():
            destination = Path(save_directory)
            for name in ("infer.py", "render.py", "setup_render.py", "requirements.txt", "requirements-render.txt",
                         "LICENSE", "THIRD_PARTY_NOTICES.md"):
                if (source_dir / name).is_file() and (source_dir / name).resolve() != (destination / name).resolve():
                    shutil.copy2(source_dir / name, destination / name)
            assets = source_dir / "render_assets"
            if (assets / "manifest.json").is_file() and assets.resolve() != (destination / "render_assets").resolve():
                from .rendering_sheetsage2 import _verify_assets
                try:
                    _verify_assets(assets, audio=True)
                except (FileNotFoundError, ValueError):
                    pass
                else:
                    shutil.copytree(assets, destination / "render_assets", dirs_exist_ok=True)
        return result

    def _prepare_audio(self, input_values, attention_mask=None):
        if self.encoder is None or not self.lora_merged:
            raise RuntimeError("Load the model with from_pretrained before inference.")
        if input_values.ndim != 2 or input_values.shape[0] < 1 or not input_values.is_floating_point():
            raise ValueError("input_values must be floating-point [batch, samples].")
        if not torch.isfinite(input_values).all():
            raise ValueError("Audio must contain only finite samples.")
        batch, samples = input_values.shape
        minimum = self.encoder.config.minimum_input_samples
        window = round(self.config.input_audio_length * self.config.sampling_rate)
        if samples < minimum:
            raise ValueError(f"Each waveform must contain at least {minimum} samples.")
        if samples > window:
            raise ValueError("Audio exceeds one model window; use transcribe for whole songs.")
        if attention_mask is None:
            lengths = torch.full((batch,), samples, dtype=torch.long, device=input_values.device)
        else:
            if attention_mask.shape != input_values.shape:
                raise ValueError("attention_mask must have the same shape as input_values.")
            mask = attention_mask.to(device=input_values.device)
            if not ((mask == 0) | (mask == 1)).all():
                raise ValueError("attention_mask must contain zeros and ones.")
            mask = mask.bool()
            lengths = mask.sum(1)
            expected = torch.arange(samples, device=input_values.device)[None] < lengths[:, None]
            if not torch.equal(mask, expected) or (lengths < minimum).any():
                raise ValueError("attention_mask must identify a nonempty right-padded waveform.")
            input_values = input_values.masked_fill(~mask, 0)
        # The encoder attends to the complete fixed window, including this silence.
        input_values = F.pad(input_values.float(), (0, window - samples))
        stride = self.encoder.config.inputs_to_logits_ratio
        if window % stride:
            input_values = F.pad(input_values, (0, stride - window % stride))
        return input_values, lengths

    def get_audio_features(self, input_values, attention_mask=None, output_hidden_states=None, return_dict=True):
        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        waveform, lengths = self._prepare_audio(input_values, attention_mask)
        mel = self.encoder.feature_extractor(waveform)
        weight_dtype = self.encoder.subsampling_module[0].convnext_layers[0].depthwise_block[1].weight.dtype
        hidden = self.encoder.subsampling_module(mel.to(dtype=weight_dtype))
        input_hidden = hidden if output_hidden_states else None
        weights = torch.softmax(self.layer_weight, dim=0)
        mixed = hidden * weights[0]
        positions = self.encoder.embed_positions(hidden)
        states = [] if output_hidden_states else None
        for weight, layer in zip(weights[1:], self.encoder.layers):
            hidden = layer(hidden, positions)
            mixed = mixed + hidden * weight
            if states is not None:
                states.append(hidden)
        memory = self.encoder_projection(mixed)
        stride = self.encoder.config.inputs_to_logits_ratio
        frame_mask = torch.arange(memory.shape[1], device=memory.device)[None] < ((lengths + stride - 1) // stride)[:, None]
        output = SheetSage2EncoderOutput(
            encoder_last_hidden_state=memory, backbone_last_hidden_state=hidden,
            mixed_hidden_state=mixed, input_hidden_state=input_hidden,
            backbone_hidden_states=tuple(states) if states is not None else None,
            feature_attention_mask=frame_mask,
        )
        return output if return_dict else output.to_tuple()

    def encode(self, audio):
        return self.get_audio_features(audio).last_hidden_state

    def _decode(self, memory, decoder_input_ids, use_cache=False, past_key_values=None, output_hidden_states=False):
        attention_mask = None if past_key_values is not None else decoder_input_ids != self.tokenizer.pad_token
        output = self.decoder(
            input_ids=decoder_input_ids, attention_mask=attention_mask,
            encoder_hidden_states=memory, encoder_attention_mask=None,
            past_key_values=past_key_values, use_cache=use_cache,
            output_hidden_states=output_hidden_states, return_dict=True,
        )
        return self.output_projection(output.last_hidden_state), output

    def decode(self, memory, decoder_input_ids, use_cache=False, past_key_values=None):
        logits, output = self._decode(memory, decoder_input_ids, use_cache, past_key_values)
        return logits, output.past_key_values

    def forward(self, input_values=None, decoder_input_ids=None, attention_mask=None,
                encoder_outputs=None, past_key_values=None, use_cache=None,
                output_hidden_states=None, return_dict=None):
        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        use_cache = self.config.use_cache if use_cache is None else use_cache
        if decoder_input_ids is None:
            raise ValueError("decoder_input_ids is required for logits; use generate to transcribe audio.")
        if encoder_outputs is None:
            if input_values is None:
                raise ValueError("Provide input_values or encoder_outputs.")
            encoder_outputs = self.get_audio_features(input_values, attention_mask, output_hidden_states)
        if torch.is_tensor(encoder_outputs):
            encoder_outputs = SheetSage2EncoderOutput(encoder_last_hidden_state=encoder_outputs)
        logits, decoded = self._decode(encoder_outputs.last_hidden_state, decoder_input_ids,
                                       use_cache, past_key_values, output_hidden_states)
        output = SheetSage2Output(
            logits=logits, past_key_values=decoded.past_key_values,
            decoder_hidden_states=decoded.hidden_states,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            backbone_last_hidden_state=encoder_outputs.backbone_last_hidden_state if output_hidden_states else None,
            mixed_hidden_state=encoder_outputs.mixed_hidden_state if output_hidden_states else None,
            input_hidden_state=encoder_outputs.input_hidden_state if output_hidden_states else None,
            backbone_hidden_states=encoder_outputs.backbone_hidden_states if output_hidden_states else None,
            feature_attention_mask=encoder_outputs.feature_attention_mask,
        )
        return_dict = self.config.use_return_dict if return_dict is None else return_dict
        return output if return_dict else output.to_tuple()

    def generate(self, input_values, **kwargs):
        """Generate grammar-constrained symbolic tokens with autoregressive caching."""
        from .generation_sheetsage2 import generate
        return generate(self, input_values, **kwargs)

    def transcribe(self, audio, output_dir=None, *, melody_only=False, **kwargs):
        """Return transcription in memory; set output_dir to also save files.

        Accepts a path, encoded audio bytes, binary stream, or waveform with
        sampling_rate. Returns ABC text, MIDI bytes, timed events, and optional
        per-window CPU tensors. With output_dir, optional tensors are saved
        instead of retained in memory. Set melody_only=True to retain both vocal
        and instrumental melodies while omitting chords from ABC and playback;
        raw predicted annotations remain available. The default keeps full
        transcription. See the model card for rendering options.
        """
        from .pipeline_sheetsage2 import transcribe
        return transcribe(self, audio, output_dir=output_dir, melody_only=melody_only, **kwargs)


SheetSage2ForConditionalGeneration = SheetSage2Model
SheetSage2Model.register_for_auto_class("AutoModel")