chore: import upstream snapshot with attribution

python/mlc_llm/model/vision/__init__.py

@@ -0,0 +1,4 @@
+"""Common `nn.Modules` used to define LLMs in this project."""
+
+from .clip_vision import CLIPVisionConfig, CLIPVisionModel
+from .image_processing import ImageProcessor

python/mlc_llm/model/vision/clip_vision.py

@@ -0,0 +1,231 @@
+"""
+Implements the CLIP Vision Encoder.
+"""
+
+import dataclasses
+import logging
+from typing import Any, Dict, Tuple  # noqa: UP035
+
+from tvm import relax
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Module, Tensor
+from tvm.relax.frontend.nn.modules import Conv2D
+from tvm.relax.frontend.nn.op import (
+    add,
+    broadcast_to,
+    concat,
+    permute_dims,
+    reshape,
+    wrap_nested,
+)
+from tvm.relax.op import arange
+
+from mlc_llm import op as op_ext
+from mlc_llm.support.config import ConfigBase
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class CLIPVisionConfig(ConfigBase):
+    """
+    Config for the vision encoder
+    """
+
+    hidden_size: int
+    image_size: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_hidden_layers: int
+    patch_size: int
+    projection_dim: int
+    vocab_size: int
+    num_channels: int = 3
+    layer_norm_eps: float = 1e-06
+    kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)  # noqa: UP006
+
+
+class CLIPVisionEmbeddings(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        self.class_embedding = nn.Parameter((self.embed_dim,))
+        self.patch_embedding = Conv2D(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+        self.position_embedding = nn.Embedding(num=self.num_positions, dim=self.embed_dim)
+
+    def forward(self, pixel_values: Tensor) -> Tensor:
+        batch_size = pixel_values.shape[0]
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+        patch_embeds = reshape(patch_embeds, shape=(batch_size, self.embed_dim, -1))
+        patch_embeds = permute_dims(
+            patch_embeds, axes=(0, 2, 1)
+        )  # shape = [batch,grid*grid,embed_dim]
+        class_embeds = broadcast_to(
+            self.class_embedding, shape=(batch_size, 1, self.embed_dim)
+        )  # shape of (batch,1,embed_dim)
+        embeddings = concat([class_embeds, patch_embeds], dim=1)
+
+        posi_ids = reshape(
+            wrap_nested(arange(0, self.num_positions, dtype="int32"), name="arange"),
+            shape=(1, -1),
+        )
+        batch_position_embedding = broadcast_to(
+            self.position_embedding(posi_ids),
+            shape=(batch_size, self.num_positions, self.embed_dim),
+        )
+        embeddings = add(embeddings, batch_position_embedding)
+        return embeddings
+
+
+def sigmoid(x: Tensor, name: str = "sigmoid") -> Tensor:
+    """Sigmoid of a Tensor
+
+    Parameters
+    ----------
+    x : Tensor
+        Input tensor to expand.
+    name : str
+        Name hint for this operator.
+
+    Returns
+    -------
+    result : Tensor
+        Sigmoid result.
+    """
+    return wrap_nested(relax.op.sigmoid(x._expr), name)
+
+
+class QuickGELU(Module):
+    def forward(self, input_tensor: Tensor) -> Tensor:
+        return input_tensor * sigmoid(input_tensor * 1.702)
+
+
+class CLIPMLP(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.activation_fn = QuickGELU()
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class CLIPAttention(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if (self.head_dim * self.num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+    ) -> Tensor:
+        d, h = self.head_dim, self.num_heads
+        b, s, _ = hidden_states.shape  # batch_size, seq_len, embed_dim
+
+        q = self.q_proj(hidden_states).reshape(b, s, h, d)
+        k = self.k_proj(hidden_states).reshape(b, s, h, d)
+        v = self.v_proj(hidden_states).reshape(b, s, h, d)
+
+        attn_output = op_ext.attention(q, k, v, None)
+        attn_output = self.out_proj(attn_output)
+        return attn_output
+
+
+class CLIPEncoderLayer(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = CLIPAttention(config)
+        self.layer_norm1 = nn.LayerNorm(normalized_shape=self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = CLIPMLP(config)
+        self.layer_norm2 = nn.LayerNorm(normalized_shape=self.embed_dim, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states = self.self_attn(hidden_states=hidden_states)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+        return outputs
+
+
+class CLIPEncoder(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+
+    def forward(self, inputs_embeds: Tensor) -> Tensor:
+        hidden_states = inputs_embeds
+        encoder_states: Tuple[Any, ...] = ()  # noqa: UP006
+        for _, encoder_layer in enumerate(self.layers):
+            encoder_states = (*encoder_states, hidden_states)
+            layer_outputs = encoder_layer(hidden_states)
+            hidden_states = layer_outputs[0]
+        encoder_states = (*encoder_states, hidden_states)
+        return encoder_states
+
+
+class CLIPVisionTransformer(Module):
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        embed_dim = config.hidden_size
+        self.embeddings = CLIPVisionEmbeddings(config)
+        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.encoder = CLIPEncoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+    def forward(self, pixel_values: Tensor) -> Tensor:
+        hidden_states = self.embeddings(pixel_values)
+        hidden_states = self.pre_layrnorm(hidden_states)
+        encoder_outputs = self.encoder(inputs_embeds=hidden_states)
+        # Apply post_layernorm to the final encoder hidden state, matching
+        # the HuggingFace CLIPVisionTransformer which returns post-normed
+        # last_hidden_state. Intermediate states remain unnormalized.
+        last_hidden_state = self.post_layernorm(encoder_outputs[-1])
+        return (*encoder_outputs[:-1], last_hidden_state)
+
+
+class CLIPVisionModel(Module):
+    no_quantization: bool = True
+
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.vision_model = CLIPVisionTransformer(config)
+
+    def forward(self, pixel_values: Tensor) -> Tensor:
+        return self.vision_model(pixel_values)[-2]

python/mlc_llm/model/vision/image_processing.py

@@ -0,0 +1,283 @@
+"""
+Implements the CLIP Image processor.
+"""
+
+from tvm import s_tir, tirx
+from tvm.relax.frontend.nn import Module, Tensor, op
+from tvm.script import tirx as T
+
+
+def _var(dtype, size=1):
+    return T.sblock_alloc_buffer((size,), dtype, scope="local")
+
+
+class ImageProcessor(Module):
+    def __init__(self):
+        super().__init__()
+
+    def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        xo, xi = sch.split(loop_x, factors=[tile[0], None])
+        yo, yi = sch.split(loop_y, factors=[tile[1], None])
+        sch.reorder(xo, yo, xi, yi)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    def resize(self, image: Tensor, params):  # image layout:NCHW
+        assert 4 == image.ndim, "image should be 4D data tensor"
+        assert 3 == image.shape[1], "image layout should be NCHW"
+
+        def get_output_image_size(image: Tensor):
+            h = image.shape[2]
+            w = image.shape[3]
+
+            if "height" in params and "width" in params:
+                return (params["height"], params["width"])
+            elif "shortest_edge" in params:
+                short = tirx.Select(w < h, w, h)
+                long = tirx.Select(w > h, w, h)
+                requested_new_short = params["shortest_edge"]
+                new_short, new_long = (
+                    tirx.Cast("int64", requested_new_short),
+                    tirx.Cast(
+                        "int64",
+                        requested_new_short
+                        * tirx.div(
+                            tirx.Cast("float32", long),
+                            tirx.Cast("float32", short),
+                        ),
+                    ),
+                )
+                ret_h = tirx.Select(w <= h, new_long, new_short)
+                ret_w = tirx.Select(w <= h, new_short, new_long)
+                return (ret_h, ret_w)
+            elif "hd_transform" in params:
+                hd_num = 4 if "hd_num" not in params else params["hd_num"]
+                pad_num = 336 if "pad_num" not in params else params["pad_num"]
+                ratio = tirx.Select(
+                    w > h,
+                    tirx.div(tirx.Cast("float32", w), tirx.Cast("float32", h)),
+                    tirx.div(tirx.Cast("float32", h), tirx.Cast("float32", w)),
+                )
+
+                scale = tirx.ceil(tirx.sqrt(tirx.Cast("float32", hd_num) * ratio))
+
+                scale = tirx.Select(
+                    (scale * tirx.ceil(tirx.div(scale, ratio))) > hd_num,
+                    scale - 1,
+                    scale,
+                )
+                scale = tirx.Cast("int64", scale)
+
+                new_w = tirx.Select(
+                    w >= h,
+                    scale * pad_num,
+                    tirx.Cast("int64", tirx.div(scale * pad_num, ratio)),
+                )
+                new_h = tirx.Select(
+                    w >= h,
+                    tirx.Cast("int64", tirx.div(new_w, ratio)),
+                    scale * pad_num,
+                )
+                return (new_h, new_w)
+            else:
+                assert False, "not supported resize parameter"
+
+        new_h, new_w = get_output_image_size(image)
+        out = op.interpolate(image, (new_h, new_w), data_layout="NCHW", mode="linear")
+        return out
+
+    def crop(self, image: Tensor, crop_size):
+        assert 4 == image.ndim, "image should be 4D data tensor"
+        assert 3 == image.shape[1], "image layout should be NCHW"
+
+        def create_crop_func(dtype):  # , top, bottom, left, right):
+            @T.prim_func(s_tir=True)
+            def crop_func(
+                image: T.handle,
+                out: T.handle,
+                top: T.int64(),
+                bottom: T.int64(),
+                left: T.int64(),
+                right: T.int64(),
+            ):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
+                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
+                out_buf = T.match_buffer(out, (n, c, bottom - top, right - left), dtype=dtype)
+                out_h = bottom - top
+                out_w = right - left
+                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(out_h, out_w):
+                            with T.sblock("crop"):
+                                if (h_idx + T.int64(top)) < h and (w_idx + T.int64(left)) < w:
+                                    T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                    T.reads(image_buf[n_idx, c_idx, h_idx + top, w_idx + left])
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[
+                                        n_idx, c_idx, h_idx + top, w_idx + left
+                                    ]
+
+            sch = s_tir.Schedule(crop_func)
+            self.apply_schedule(sch, sch.get_sblock("crop"))
+            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)
+
+        n, c, orig_height, orig_width = image.shape
+        crop_height = crop_size["height"]
+        crop_width = crop_size["width"]
+
+        top = (orig_height - crop_height) // 2
+        bottom = orig_height - top
+
+        left = (orig_width - crop_width) // 2
+        right = orig_width - left
+
+        out = op.tensor_ir_op(
+            create_crop_func(image.dtype),
+            "crop",
+            [image, top, bottom, left, right],
+            [Tensor.placeholder([n, c, crop_height, crop_width], image.dtype)],
+        )
+        return out
+
+    def rescale(self, image: Tensor, rescale_factor=1 / 255.0, o_dtype="float32"):
+        assert 4 == image.ndim, "image should be 4D data tensor"
+        assert 3 == image.shape[1], "image layout should be NCHW"
+
+        def create_rescale_func(rescale_factor, dtype, o_dtype):
+            @T.prim_func(s_tir=True)
+            def rescale_func(image: T.handle, out: T.handle):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
+                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
+                out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype)
+
+                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(h, w):
+                            with T.sblock("rescale"):
+                                T.reads(image_buf[n_idx, c_idx, h_idx, w_idx])
+                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                if h_idx < h and w_idx < w:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = (
+                                        T.cast(
+                                            image_buf[n_idx, c_idx, h_idx, w_idx],
+                                            o_dtype,
+                                        )
+                                        * rescale_factor
+                                    )
+
+            sch = s_tir.Schedule(rescale_func)
+            self.apply_schedule(sch, sch.get_sblock("rescale"))
+            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)
+
+        out = op.tensor_ir_op(
+            create_rescale_func(rescale_factor, image.dtype, o_dtype),
+            "rescale",
+            [image],
+            [Tensor.placeholder(image.shape, o_dtype)],
+        )
+        return out
+
+    def normalize(self, image: Tensor, o_dtype="float32"):
+        assert 4 == image.ndim, "image should be 4D data tensor"
+        assert 3 == image.shape[1], "image layout should be NCHW"
+
+        def create_normalize_func(dtype, o_dtype):
+            @T.prim_func(s_tir=True)
+            def normalize_func(image: T.handle, out: T.handle):
+                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
+                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
+                out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype)
+                mean = _var(o_dtype, 3)
+                stddev = _var(o_dtype, 3)
+
+                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(h, w):
+                            with T.sblock("normalize"):
+                                T.reads(
+                                    image_buf[n_idx, c_idx, h_idx, w_idx],
+                                    mean[c_idx],
+                                    stddev[c_idx],
+                                )
+                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                with T.init():
+                                    mean[0] = 0.48145466
+                                    stddev[0] = 0.26862954
+                                    mean[1] = 0.4578275
+                                    stddev[1] = 0.26130258
+                                    mean[2] = 0.40821073
+                                    stddev[2] = 0.27577711
+                                if h_idx < h and w_idx < w:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = (
+                                        T.cast(
+                                            image_buf[n_idx, c_idx, h_idx, w_idx],
+                                            o_dtype,
+                                        )
+                                        - mean[c_idx]
+                                    ) / stddev[c_idx]
+
+            sch = s_tir.Schedule(normalize_func)
+            self.apply_schedule(sch, sch.get_sblock("normalize"))
+            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)
+
+        out = op.tensor_ir_op(
+            create_normalize_func(image.dtype, o_dtype),
+            "normalize",
+            [image],
+            [Tensor.placeholder(image.shape, o_dtype)],
+        )
+        return out
+
+    def pad(self, image: Tensor, dtype="uint8"):
+        assert 4 == image.ndim, "image should be 4D data tensor"
+        assert 3 == image.shape[1], "image layout should be NCHW"
+
+        def create_pad_func(left, right, fill=255):
+            @T.prim_func(s_tir=True)
+            def pad_func(image: T.handle, out: T.handle, t: T.int64(), b: T.int64()):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
+                image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype)
+                out_buf = T.match_buffer(out, (n, c, h + t + b, w + left + right), dtype=dtype)
+                out_h = h + t + b
+                out_w = w + left + right
+
+                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(out_h, out_w):
+                            with T.sblock("pad"):
+                                T.reads(image_buf[n_idx, c_idx, h_idx, w_idx])
+                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                if h_idx < t or h_idx > h + b or w_idx < left or w_idx > w + right:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = fill
+                                else:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[
+                                        n_idx, c_idx, h_idx - t, w_idx - left
+                                    ]
+
+            sch = s_tir.Schedule(pad_func)
+            self.apply_schedule(sch, sch.get_sblock("pad"))
+            return sch.mod["main"].with_attr("tirx.is_scheduled", 1)
+
+        h = image.shape[2]
+        tar = tirx.truncdiv(h + 335, 336) * 336
+        t = tirx.div(tar - h, 2)
+        b = tar - h - t
+        left = 0
+        right = 0
+
+        n, c, h, w = image.shape
+        out = op.tensor_ir_op(
+            create_pad_func(left, right),
+            "pad",
+            [image, t, b],
+            [Tensor.placeholder((n, c, tar, w), image.dtype)],
+        )
+        return out
+
+    def preprocess(self, pixel_values):
+        return pixel_values