Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import MmprojModel, ModelBase, gguf11 12 13@ModelBase.register("KimiVLForConditionalGeneration")14class KimiVLModel(MmprojModel):15    def __init__(self, *args, **kwargs):16        super().__init__(*args, **kwargs)17        assert self.hparams_vision is not None18        self.hparams_vision["image_size"] = 64 * 14 # for compatibility19 20    def set_gguf_parameters(self):21        super().set_gguf_parameters()22        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.KIMIVL)23        self.gguf_writer.add_vision_use_gelu(True)24        self.gguf_writer.add_vision_projector_scale_factor(2)25        # eps is the same as pytorch's default value26        assert self.hparams_vision is not None27        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-5))28 29    @classmethod30    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:31        name, gen = item32 33        is_vision_tensor = "vision_tower" in name or "multi_modal_projector" in name34 35        if not is_vision_tensor:36            return None37 38        return super().filter_tensors(item)39 40    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:41        if "pos_emb.weight" in name:42            data_torch = data_torch.view(data_torch.shape[0] * data_torch.shape[1], data_torch.shape[2])43 44        if "wqkv" in name:45            split_dim = 0 if "weight" in name else -146            wq, wk, wv = data_torch.chunk(3, dim=split_dim)47            yield from super().modify_tensors(wq, name.replace("wqkv", "wq"), bid)48            yield from super().modify_tensors(wk, name.replace("wqkv", "wk"), bid)49            yield from super().modify_tensors(wv, name.replace("wqkv", "wv"), bid)50        else:51            yield from super().modify_tensors(data_torch, name, bid)52 53 54@ModelBase.register("KimiK25ForConditionalGeneration")55class KimiK25Model(MmprojModel):56    """Kimi-K2.5 with MoonViT3d vision encoder"""57 58    def __init__(self, *args, **kwargs):59        super().__init__(*args, **kwargs)60 61        assert self.hparams_vision is not None, "Kimi-K2.5 requires vision_config in model config"62 63        self.merge_kernel_size = tuple(self.hparams_vision.get("merge_kernel_size", [2, 2]))64        self.patch_size = self.hparams_vision.get("patch_size", 14)65 66        # Set image_size for compatibility with base class67        # Use position embedding dimensions as image_size reference68        pos_emb_h = self.hparams_vision.get("init_pos_emb_height", 64)69        self.hparams_vision["image_size"] = pos_emb_h * self.patch_size70 71    def set_gguf_parameters(self):72        # Base class MmprojModel.set_gguf_parameters() already writes:73        # - vision_block_count, vision_head_count, vision_embedding_length74        # - vision_feed_forward_length, vision_patch_size, image_mean, image_std75        # via find_vparam() which handles the vt_* prefixed keys in Kimi-K2.5's config76        super().set_gguf_parameters()77        assert self.hparams_vision is not None78 79        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.KIMIK25)80 81        # Position embedding parameters (for interpolation)82        self.gguf_writer.add_uint32("vision.pos_emb_height", self.hparams_vision.get("init_pos_emb_height", 64))83        self.gguf_writer.add_uint32("vision.pos_emb_width", self.hparams_vision.get("init_pos_emb_width", 64))84        self.gguf_writer.add_uint32("vision.pos_emb_time", self.hparams_vision.get("init_pos_emb_time", 4))85 86        # Projector parameters87        self.gguf_writer.add_vision_use_gelu(self.hparams_vision.get("projector_hidden_act", "gelu") == "gelu")88        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("projector_ln_eps", 1e-5))89        self.gguf_writer.add_vision_projector_scale_factor(self.merge_kernel_size[0])90 91        # Image size limits92        # Note: in_patch_limit is for images, in_patch_limit_each_frame is for video (not supported yet)93        in_patch_limit = self.preprocessor_config.get("in_patch_limit", 16384)94        min_patches = 8  # reasonable minimum95        pixels_per_patch = self.patch_size ** 296        self.gguf_writer.add_vision_min_pixels(min_patches * pixels_per_patch)97        self.gguf_writer.add_vision_max_pixels(in_patch_limit * pixels_per_patch)98 99    @staticmethod100    def permute(weights: Tensor, n_head: int) -> Tensor:101        out_dim, in_dim = weights.shape102        head_dim = out_dim // n_head103        w = weights.reshape(n_head, head_dim // 4, 2, 2, in_dim)104        w = w.permute(0, 2, 1, 3, 4)105        return w.reshape(out_dim, in_dim)106 107    @classmethod108    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:109        name, gen = item110 111        # Only process vision and projector tensors112        is_vision = any(x in name for x in ["vision_tower", "mm_projector"])113 114        if not is_vision:115            return None116 117        return super().filter_tensors(item)118 119    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:120        assert self.hparams_vision is not None121        n_head = self.hparams_vision.get("num_attention_heads", 16)122 123        # Permute Q/K weights/biases from interleaved to split RoPE format124        # This allows using build_rope_2d at runtime without post-permutation.125        if "wqkv" in name:126            out_dim = data_torch.shape[0]127            qkv_dim = out_dim // 3128            head_dim = qkv_dim // n_head129 130            if "weight" in name:131                wq, wk, wv = data_torch[:qkv_dim, :], data_torch[qkv_dim:2 * qkv_dim, :], data_torch[2 * qkv_dim:, :]132                wq = self.permute(wq, n_head)133                wk = self.permute(wk, n_head)134                data_torch = torch.cat([wq, wk, wv], dim=0)135            elif "bias" in name:136                bq, bk, bv = data_torch[:qkv_dim], data_torch[qkv_dim:2 * qkv_dim], data_torch[2 * qkv_dim:]137                bq = bq.reshape(n_head, head_dim // 4, 2, 2).permute(0, 2, 1, 3).reshape(-1)138                bk = bk.reshape(n_head, head_dim // 4, 2, 2).permute(0, 2, 1, 3).reshape(-1)139                data_torch = torch.cat([bq, bk, bv], dim=0)140 141        # Temporal embeddings: (T, 1, C) → (T, C)142        if "pos_emb.time_weight" in name:143            T, _, C = data_torch.shape144            data_torch = data_torch.reshape(T, C)145 146        # PatchMergerMLP tensor name mapping147        # proj.0.weight → proj.linear_1.weight148        # proj.2.weight → proj.linear_2.weight149        if "mm_projector.proj.0." in name:150            name = name.replace(".proj.0.", ".proj.linear_1.")151        elif "mm_projector.proj.2." in name:152            name = name.replace(".proj.2.", ".proj.linear_2.")153 154        yield from super().modify_tensors(data_torch, name, bid)155 156 157@ModelBase.register("Glm5vForConditionalGeneration")158class Glm5vModel(KimiK25Model):159    """GLM-5.2-Vision MoonViT3d encoder and projector160 161    Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the162    kimik25 projector type. The image begin/end tokens differ, but they are163    resolved at runtime from the text model vocab.164    """165 166    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:167        if name.startswith("mm_projector.linear_"):168            name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)169 170        yield from super().modify_tensors(data_torch, name, bid)171