Felipe97/llama-cpp-compiled

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")14@ModelBase.example("moonshotai/Kimi-VL-A3B-Instruct")15class KimiVLModel(MmprojModel):16    def __init__(self, *args, **kwargs):17        super().__init__(*args, **kwargs)18        assert self.hparams_vision is not None19        self.hparams_vision["image_size"] = 64 * 14 # for compatibility20 21    def set_gguf_parameters(self):22        super().set_gguf_parameters()23        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.KIMIVL)24        self.gguf_writer.add_vision_use_gelu(True)25        self.gguf_writer.add_vision_projector_scale_factor(2)26        # eps is the same as pytorch's default value27        assert self.hparams_vision is not None28        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-5))29 30    @classmethod31    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:32        name, gen = item33 34        is_vision_tensor = "vision_tower" in name or "multi_modal_projector" in name35 36        if not is_vision_tensor:37            return None38 39        return super().filter_tensors(item)40 41    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:42        if "pos_emb.weight" in name:43            data_torch = data_torch.view(data_torch.shape[0] * data_torch.shape[1], data_torch.shape[2])44 45        if "wqkv" in name:46            split_dim = 0 if "weight" in name else -147            wq, wk, wv = data_torch.chunk(3, dim=split_dim)48            yield from super().modify_tensors(wq, name.replace("wqkv", "wq"), bid)49            yield from super().modify_tensors(wk, name.replace("wqkv", "wk"), bid)50            yield from super().modify_tensors(wv, name.replace("wqkv", "wv"), bid)51        else:52            yield from super().modify_tensors(data_torch, name, bid)53 54 55@ModelBase.register("KimiK25ForConditionalGeneration")56@ModelBase.example("moonshotai/Kimi-K2.5")57class KimiK25Model(MmprojModel):58    """Kimi-K2.5 with MoonViT3d vision encoder"""59 60    def __init__(self, *args, **kwargs):61        super().__init__(*args, **kwargs)62 63        assert self.hparams_vision is not None, "Kimi-K2.5 requires vision_config in model config"64 65        self.merge_kernel_size = tuple(self.hparams_vision.get("merge_kernel_size", [2, 2]))66        self.patch_size = self.hparams_vision.get("patch_size", 14)67 68        # Set image_size for compatibility with base class69        # Use position embedding dimensions as image_size reference70        pos_emb_h = self.hparams_vision.get("init_pos_emb_height", 64)71        self.hparams_vision["image_size"] = pos_emb_h * self.patch_size72 73    def set_gguf_parameters(self):74        # Base class MmprojModel.set_gguf_parameters() already writes:75        # - vision_block_count, vision_head_count, vision_embedding_length76        # - vision_feed_forward_length, vision_patch_size, image_mean, image_std77        # via find_vparam() which handles the vt_* prefixed keys in Kimi-K2.5's config78        super().set_gguf_parameters()79        assert self.hparams_vision is not None80 81        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.KIMIK25)82 83        # Position embedding parameters (for interpolation)84        self.gguf_writer.add_uint32("vision.pos_emb_height", self.hparams_vision.get("init_pos_emb_height", 64))85        self.gguf_writer.add_uint32("vision.pos_emb_width", self.hparams_vision.get("init_pos_emb_width", 64))86        self.gguf_writer.add_uint32("vision.pos_emb_time", self.hparams_vision.get("init_pos_emb_time", 4))87 88        # Projector parameters89        self.gguf_writer.add_vision_use_gelu(self.hparams_vision.get("projector_hidden_act", "gelu") == "gelu")90        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("projector_ln_eps", 1e-5))91        self.gguf_writer.add_vision_projector_scale_factor(self.merge_kernel_size[0])92 93        # Image size limits94        # Note: in_patch_limit is for images, in_patch_limit_each_frame is for video (not supported yet)95        in_patch_limit = self.preprocessor_config.get("in_patch_limit", 16384)96        min_patches = 8  # reasonable minimum97        pixels_per_patch = self.patch_size ** 298        self.gguf_writer.add_vision_min_pixels(min_patches * pixels_per_patch)99        self.gguf_writer.add_vision_max_pixels(in_patch_limit * pixels_per_patch)100 101    @staticmethod102    def permute(weights: Tensor, n_head: int) -> Tensor:103        out_dim, in_dim = weights.shape104        head_dim = out_dim // n_head105        w = weights.reshape(n_head, head_dim // 4, 2, 2, in_dim)106        w = w.permute(0, 2, 1, 3, 4)107        return w.reshape(out_dim, in_dim)108 109    @classmethod110    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:111        name, gen = item112 113        # Only process vision and projector tensors114        is_vision = any(x in name for x in ["vision_tower", "mm_projector"])115 116        if not is_vision:117            return None118 119        return super().filter_tensors(item)120 121    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:122        assert self.hparams_vision is not None123        n_head = self.hparams_vision.get("num_attention_heads", 16)124 125        # Permute Q/K weights/biases from interleaved to split RoPE format126        # This allows using build_rope_2d at runtime without post-permutation.127        if "wqkv" in name:128            out_dim = data_torch.shape[0]129            qkv_dim = out_dim // 3130            head_dim = qkv_dim // n_head131 132            if "weight" in name:133                wq, wk, wv = data_torch[:qkv_dim, :], data_torch[qkv_dim:2 * qkv_dim, :], data_torch[2 * qkv_dim:, :]134                wq = self.permute(wq, n_head)135                wk = self.permute(wk, n_head)136                data_torch = torch.cat([wq, wk, wv], dim=0)137            elif "bias" in name:138                bq, bk, bv = data_torch[:qkv_dim], data_torch[qkv_dim:2 * qkv_dim], data_torch[2 * qkv_dim:]139                bq = bq.reshape(n_head, head_dim // 4, 2, 2).permute(0, 2, 1, 3).reshape(-1)140                bk = bk.reshape(n_head, head_dim // 4, 2, 2).permute(0, 2, 1, 3).reshape(-1)141                data_torch = torch.cat([bq, bk, bv], dim=0)142 143        # Temporal embeddings: (T, 1, C) → (T, C)144        if "pos_emb.time_weight" in name:145            T, _, C = data_torch.shape146            data_torch = data_torch.reshape(T, C)147 148        # PatchMergerMLP tensor name mapping149        # proj.0.weight → proj.linear_1.weight150        # proj.2.weight → proj.linear_2.weight151        if "mm_projector.proj.0." in name:152            name = name.replace(".proj.0.", ".proj.linear_1.")153        elif "mm_projector.proj.2." in name:154            name = name.replace(".proj.2.", ".proj.linear_2.")155 156        yield from super().modify_tensors(data_torch, name, bid)157 158 159@ModelBase.register("Glm5vForConditionalGeneration")160# [TAG_HF_EXAMPLE_MISSING]161class Glm5vModel(KimiK25Model):162    """GLM-5.2-Vision MoonViT3d encoder and projector163 164    Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the165    kimik25 projector type. The image begin/end tokens differ, but they are166    resolved at runtime from the text model vocab.167    """168 169    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:170        if name.startswith("mm_projector.linear_"):171            name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)172 173        yield from super().modify_tensors(data_torch, name, bid)174