Upload folder using huggingface_hub

Browse files

Files changed (3) hide show

README.md +2 -3
build/torch-universal/kernel_image_resize/__init__.py +2 -1
build/torch-universal/kernel_image_resize/_ragged.py +217 -0

README.md CHANGED Viewed

@@ -42,9 +42,8 @@ pixel_values = kir.resize_normalize(
 # -> (3, 3, 384, 384) float32, ready for the model
 ```
-`trust_remote_code=True` is required because this is a personal namespace (not the trusted
-`kernels-community` org). `revision="main"` loads the current code; tag a `v1.0.0` release if
-you want `version=1` loading instead.
 `resize_normalize` accepts a stacked `(N, C, H, W)` tensor or a ragged list of CHW
 tensors. `resize_normalize_ragged` is the same kernel, list-only.

 # -> (3, 3, 384, 384) float32, ready for the model
 ```
+Requires `kernels >= 0.15` (published as a `kernel` repo type). `trust_remote_code=True` is needed
+because `Molbap` is a personal namespace, not the auto-trusted `kernels-community` org.
 `resize_normalize` accepts a stacked `(N, C, H, W)` tensor or a ragged list of CHW
 tensors. `resize_normalize_ragged` is the same kernel, list-only.

build/torch-universal/kernel_image_resize/__init__.py CHANGED Viewed

@@ -18,6 +18,7 @@ launch, taps*taps. Both parity <=1e-4 vs torchvision-float.
 from ._fused import fused_resize_normalize
 from ._pack import PIL_RESAMPLE_TO_INTERP, as_image_list
 from ._separable import separable_resize_crop_normalize, separable_resize_normalize
@@ -110,4 +111,4 @@ def resize_normalize_ragged(
     )
-__all__ = ["resize_normalize", "resize_normalize_ragged"]

 from ._fused import fused_resize_normalize
 from ._pack import PIL_RESAMPLE_TO_INTERP, as_image_list
+from ._ragged import resize_normalize_ragged_output
 from ._separable import separable_resize_crop_normalize, separable_resize_normalize
     )
+__all__ = ["resize_normalize", "resize_normalize_ragged", "resize_normalize_ragged_output"]

build/torch-universal/kernel_image_resize/_ragged.py ADDED Viewed

	@@ -0,0 +1,217 @@

+"""Ragged-OUTPUT resize + normalize, for dynamic-resolution VLMs.
+The classic kernel (`resize_normalize`) takes ragged input and produces ONE fixed output size.
+Recent VLMs do the opposite at the output end: each image is resized to its OWN target size
+(aspect-preserving, bounded by a patch budget) and then patchified — Gemma4, Qwen-style
+`smart_resize` (minimax_m3_vl, glmga), etc. So the output is ragged: image i -> (C, h_i, w_i).
+This kernel keeps the classic kernel's efficiency (two separable passes, uint8 input, taps+taps,
+TWO launches for the whole batch — no per-image Python loop) but writes into a packed ragged
+output buffer with per-image offsets, and returns a list of per-image (C, h_i, w_i) tensors.
+The caller computes the per-image target sizes with the model's own rule (Gemma4's
+`get_aspect_ratio_preserving_size`, Qwen's `smart_resize`, ...) and passes them in. Patchify is
+left to the model (it is a cheap reshape/unfold over this kernel's output).
+EXPERIMENTAL: authored without a GPU; validate parity on CUDA before use.
+"""
+import triton
+import triton.language as tl
+from ._fused import _resample_weight
+from ._pack import fold_mean_std, pack_images
+@triton.jit
+def _ragged_horizontal_kernel(
+    input_pixels, intermediate, input_offsets, intermediate_offsets,
+    heights, widths, out_widths,
+    cubic_coeff,
+    CHANNELS: tl.constexpr, BLOCK: tl.constexpr, CUBIC: tl.constexpr, ANTIALIAS: tl.constexpr,
+    MAX_TAPS_COL: tl.constexpr,
+):
+    """Resize width per image: uint8 (C, H, W) -> float (C, H, w_i)."""
+    image_index = tl.program_id(0)
+    block_index = tl.program_id(1)
+    in_height = tl.load(heights + image_index)
+    in_width = tl.load(widths + image_index)
+    out_width = tl.load(out_widths + image_index)
+    input_start = tl.load(input_offsets + image_index)
+    intermediate_start = tl.load(intermediate_offsets + image_index)
+    in_width_f = in_width.to(tl.float32)
+    num_pixels = in_height * out_width
+    flat_index = block_index * BLOCK + tl.arange(0, BLOCK)
+    active = flat_index < num_pixels
+    input_row = flat_index // out_width
+    out_col = flat_index % out_width
+    filter_half = 2.0 if CUBIC else 1.0
+    col_scale = in_width_f / out_width.to(tl.float32)
+    col_filter_scale = tl.maximum(col_scale, 1.0) if ANTIALIAS else 1.0
+    col_support = filter_half * col_filter_scale
+    col_inv_scale = 1.0 / col_filter_scale
+    src_center_col = col_scale * (out_col.to(tl.float32) + 0.5)
+    first_tap_col = tl.floor(src_center_col - col_support + 0.5)
+    col_weight_sum = tl.zeros([BLOCK], dtype=tl.float32)
+    for tap in tl.static_range(MAX_TAPS_COL):
+        tap_col = first_tap_col + tap
+        weight = _resample_weight((tap_col - src_center_col + 0.5) * col_inv_scale, cubic_coeff, CUBIC)
+        if ANTIALIAS:
+            weight = tl.where((tap_col >= 0.0) & (tap_col < in_width_f), weight, 0.0)
+        col_weight_sum += weight
+    input_plane = (in_height * in_width).to(tl.int64)
+    intermediate_plane = (in_height * out_width).to(tl.int64)
+    in_width_i64 = in_width.to(tl.int64)
+    out_width_i64 = out_width.to(tl.int64)
+    input_row_i64 = input_row.to(tl.int64)
+    for channel in tl.static_range(CHANNELS):
+        input_row_base = input_start + channel * input_plane + input_row_i64 * in_width_i64
+        accumulator = tl.zeros([BLOCK], dtype=tl.float32)
+        for tap in tl.static_range(MAX_TAPS_COL):
+            tap_col = first_tap_col + tap
+            weight = _resample_weight((tap_col - src_center_col + 0.5) * col_inv_scale, cubic_coeff, CUBIC)
+            if ANTIALIAS:
+                weight = tl.where((tap_col >= 0.0) & (tap_col < in_width_f), weight, 0.0)
+            clamped_tap_col = tl.minimum(tl.maximum(tap_col.to(tl.int32), 0), in_width - 1).to(tl.int64)
+            pixel = tl.load(input_pixels + input_row_base + clamped_tap_col, mask=active, other=0).to(tl.float32)
+            accumulator += weight * pixel
+        accumulator = accumulator / col_weight_sum
+        write_index = intermediate_start + channel * intermediate_plane + input_row_i64 * out_width_i64 + out_col
+        tl.store(intermediate + write_index, accumulator, mask=active)
+@triton.jit
+def _ragged_vertical_kernel(
+    intermediate, output, intermediate_offsets, output_offsets,
+    heights, out_heights, out_widths, means, stds,
+    cubic_coeff,
+    CHANNELS: tl.constexpr, BLOCK: tl.constexpr, CUBIC: tl.constexpr, ANTIALIAS: tl.constexpr,
+    MAX_TAPS_ROW: tl.constexpr,
+):
+    """Resize height per image + normalize: float (C, H, w_i) -> float (C, h_i, w_i)."""
+    image_index = tl.program_id(0)
+    block_index = tl.program_id(1)
+    in_height = tl.load(heights + image_index)
+    out_height = tl.load(out_heights + image_index)
+    out_width = tl.load(out_widths + image_index)
+    intermediate_start = tl.load(intermediate_offsets + image_index)
+    output_start = tl.load(output_offsets + image_index)
+    in_height_f = in_height.to(tl.float32)
+    num_pixels = out_height * out_width
+    flat_index = block_index * BLOCK + tl.arange(0, BLOCK)
+    active = flat_index < num_pixels
+    out_row = flat_index // out_width
+    out_col = flat_index % out_width
+    filter_half = 2.0 if CUBIC else 1.0
+    row_scale = in_height_f / out_height.to(tl.float32)
+    row_filter_scale = tl.maximum(row_scale, 1.0) if ANTIALIAS else 1.0
+    row_support = filter_half * row_filter_scale
+    row_inv_scale = 1.0 / row_filter_scale
+    src_center_row = row_scale * (out_row.to(tl.float32) + 0.5)
+    first_tap_row = tl.floor(src_center_row - row_support + 0.5)
+    row_weight_sum = tl.zeros([BLOCK], dtype=tl.float32)
+    for tap in tl.static_range(MAX_TAPS_ROW):
+        tap_row = first_tap_row + tap
+        weight = _resample_weight((tap_row - src_center_row + 0.5) * row_inv_scale, cubic_coeff, CUBIC)
+        if ANTIALIAS:
+            weight = tl.where((tap_row >= 0.0) & (tap_row < in_height_f), weight, 0.0)
+        row_weight_sum += weight
+    intermediate_plane = (in_height * out_width).to(tl.int64)
+    output_plane = (out_height * out_width).to(tl.int64)
+    out_width_i64 = out_width.to(tl.int64)
+    out_col_i64 = out_col.to(tl.int64)
+    out_row_i64 = out_row.to(tl.int64)
+    for channel in tl.static_range(CHANNELS):
+        channel_base = intermediate_start + channel * intermediate_plane
+        accumulator = tl.zeros([BLOCK], dtype=tl.float32)
+        for tap in tl.static_range(MAX_TAPS_ROW):
+            tap_row = first_tap_row + tap
+            weight = _resample_weight((tap_row - src_center_row + 0.5) * row_inv_scale, cubic_coeff, CUBIC)
+            if ANTIALIAS:
+                weight = tl.where((tap_row >= 0.0) & (tap_row < in_height_f), weight, 0.0)
+            clamped_tap_row = tl.minimum(tl.maximum(tap_row.to(tl.int32), 0), in_height - 1).to(tl.int64)
+            pixel = tl.load(intermediate + channel_base + clamped_tap_row * out_width_i64 + out_col_i64, mask=active, other=0.0)
+            accumulator += weight * pixel
+        accumulator = accumulator / row_weight_sum
+        mean = tl.load(means + channel)
+        std = tl.load(stds + channel)
+        accumulator = (accumulator - mean) / std
+        write_index = output_start + channel * output_plane + out_row_i64 * out_width_i64 + out_col_i64
+        tl.store(output + write_index, accumulator, mask=active)
+def _ragged_max_taps(in_sizes, out_sizes, interp, antialias):
+    import math
+    interp_half = 2.0 if interp == "bicubic" else 1.0
+    worst = 0
+    for in_size, out_size in zip(in_sizes, out_sizes):
+        scale = in_size / out_size
+        eff = max(scale, 1.0) if antialias else 1.0
+        worst = max(worst, math.ceil(interp_half * eff) * 2 + 1)
+    return worst
+def resize_normalize_ragged_output(images, target_sizes, mean, std, rescale, interp, antialias, block: int = 256):
+    """Resize each image to its own `target_sizes[i] = (h_i, w_i)` and normalize.
+    Returns a list of per-image `(C, h_i, w_i)` float tensors (views into one packed buffer).
+    `target_sizes` come from the model's resize rule (e.g. Gemma4 aspect-ratio-preserving,
+    Qwen smart_resize). No cropping — VLMs resize then patchify.
+    """
+    import torch
+    images = list(images)
+    device = images[0].device
+    num_images = len(images)
+    cubic_coeff = -0.5 if antialias else -0.75
+    in_heights = [int(img.shape[1]) for img in images]
+    in_widths = [int(img.shape[2]) for img in images]
+    out_heights = [int(h) for h, _ in target_sizes]
+    out_widths = [int(w) for _, w in target_sizes]
+    max_taps_row = _ragged_max_taps(in_heights, out_heights, interp, antialias)
+    max_taps_col = _ragged_max_taps(in_widths, out_widths, interp, antialias)
+    means, stds = fold_mean_std(mean, std, rescale, device)
+    input_pixels, input_offsets, heights, widths, channels = pack_images(images, dtype=torch.uint8)
+    intermediate_offsets_list, intermediate_total, max_mid = [], 0, 0
+    output_offsets_list, output_total, max_out = [], 0, 0
+    for h_in, w_out, h_out in zip(in_heights, out_widths, out_heights):
+        intermediate_offsets_list.append(intermediate_total)
+        intermediate_total += channels * h_in * w_out
+        max_mid = max(max_mid, h_in * w_out)
+        output_offsets_list.append(output_total)
+        output_total += channels * h_out * w_out
+        max_out = max(max_out, h_out * w_out)
+    intermediate = torch.empty(intermediate_total, device=device, dtype=torch.float32)
+    output = torch.empty(output_total, device=device, dtype=torch.float32)
+    intermediate_offsets = torch.tensor(intermediate_offsets_list, device=device, dtype=torch.int64)
+    output_offsets = torch.tensor(output_offsets_list, device=device, dtype=torch.int64)
+    out_widths_t = torch.tensor(out_widths, device=device, dtype=torch.int32)
+    out_heights_t = torch.tensor(out_heights, device=device, dtype=torch.int32)
+    _ragged_horizontal_kernel[(num_images, triton.cdiv(max_mid, block))](
+        input_pixels, intermediate, input_offsets, intermediate_offsets, heights, widths, out_widths_t,
+        cubic_coeff, CHANNELS=channels, BLOCK=block, CUBIC=(interp == "bicubic"), ANTIALIAS=antialias,
+        MAX_TAPS_COL=max_taps_col,
+    )
+    _ragged_vertical_kernel[(num_images, triton.cdiv(max_out, block))](
+        intermediate, output, intermediate_offsets, output_offsets, heights, out_heights_t, out_widths_t,
+        means, stds, cubic_coeff, CHANNELS=channels, BLOCK=block, CUBIC=(interp == "bicubic"),
+        ANTIALIAS=antialias, MAX_TAPS_ROW=max_taps_row,
+    )
+    results = []
+    for offset, h_out, w_out in zip(output_offsets_list, out_heights, out_widths):
+        results.append(output[offset : offset + channels * h_out * w_out].view(channels, h_out, w_out))
+    return results