Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import MmprojModel, ModelBase, gguf9 10from .llama import LlamaModel11 12 13@ModelBase.register("JanusForConditionalGeneration")14class JanusProModel(LlamaModel):15    model_arch = gguf.MODEL_ARCH.LLAMA  # reuse Llama arch16 17    @classmethod18    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:19        name, gen = item20 21        # Skip vision, aligner, and generation tensors22        skip_prefixes = (23            'model.vision_model.',24            'model.aligner.',25            'model.vqmodel.',26            'model.generation_embeddings.',27            'model.generation_aligner.',28            'model.generation_head.',29        )30        if name.startswith(skip_prefixes):31            return None32 33        return super().filter_tensors(item)34 35 36@ModelBase.register("JanusForConditionalGeneration")37class JanusProVisionModel(MmprojModel):38    def __init__(self, *args, **kwargs):39        super().__init__(*args, **kwargs)40        assert self.hparams_vision is not None41        if "intermediate_size" not in self.hparams_vision:42            mlp_ratio = self.hparams_vision.get("mlp_ratio")43            hidden_size = self.hparams_vision.get("hidden_size")44            if mlp_ratio is not None and hidden_size is not None:45                self.hparams_vision["intermediate_size"] = int(round(hidden_size * mlp_ratio))46 47    def set_gguf_parameters(self):48        super().set_gguf_parameters()49        assert self.hparams_vision is not None50 51        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.JANUS_PRO)52 53        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-6))54 55        hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower()56        if hidden_act == "gelu":57            self.gguf_writer.add_vision_use_gelu(True)58        elif hidden_act == "silu":59            self.gguf_writer.add_vision_use_silu(True)60 61    def _map_aligner_tensor(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]:62        """Map aligner tensors to projector format"""63        suffix = ".bias" if name.endswith(".bias") else ".weight"64 65        if name.startswith("model.aligner."):66            local_name = name[len("model.aligner."):]67        elif name.startswith("aligner."):68            local_name = name[len("aligner."):]69        else:70            raise ValueError(f"Unsupported Janus aligner prefix: {name}")71 72        if local_name.startswith("fc1."):73            mm_index = 074        elif local_name.startswith("hidden_layers."):75            parts = local_name.split(".", 2)76            if len(parts) < 3:77                raise ValueError(f"Unexpected Janus aligner tensor name: {name}")78            mm_index = int(parts[1]) + 179        else:80            raise ValueError(f"Unsupported Janus aligner tensor: {name}")81 82        tensor_name = self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, mm_index, suffix=suffix)83        return [(tensor_name, data_torch)]84 85    @classmethod86    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:87        name, gen = item88 89        # Skip generation-related components90        skip_generation_prefixes = (91            'model.vqmodel.',92            'vqmodel.',93            'model.generation_embeddings.',94            'generation_embeddings.',95            'model.generation_aligner.',96            'generation_aligner.',97            'model.generation_head.',98            'generation_head.',99        )100        if name.startswith(skip_generation_prefixes):101            return None102 103        return super().filter_tensors(item)104 105    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:106        # Handle aligner tensors107        if name.startswith(('model.aligner.', 'aligner.')):108            yield from self._map_aligner_tensor(data_torch, name)109            return110 111        # Handle vision tensors112        if name.startswith(('model.vision_model.', 'vision_model.')):113            yield from super().modify_tensors(data_torch, name, bid)114            return115 116        return117