llmcompressor.modeling.kimi_k3.media_utils

def fill_transparent_bg_with(
    image: Image.Image,
    transparent_bg_config: TransparentBgConfig | None = None,
) -> Image.Image:
    """Composite a (possibly) transparent image onto a configured background.

    When ``transparent_bg_config`` is ``None``, the image is simply converted
    to RGB (preserving the historical behavior). Otherwise the alpha channel
    is alpha-composited over a background generated according to the config.
    """
    if transparent_bg_config is None:
        return image.convert("RGB")

    if image.mode == "RGB":
        return image

    has_alpha = "A" in image.getbands() or "transparency" in image.info
    if not has_alpha:
        return image.convert("RGB")

    img = np.array(image.convert("RGBA"))
    height, width = img.shape[:2]
    bg_pattern = transparent_bg_config.pattern
    if bg_pattern == "white":
        bg = np.full((height, width, 3), 255, dtype=np.uint8)
    elif bg_pattern == "black":
        bg = np.zeros((height, width, 3), dtype=np.uint8)
    elif bg_pattern == "gray":
        bg = np.full((height, width, 3), 128, dtype=np.uint8)
    elif bg_pattern == "chessboard":
        bg = _create_chessboard_background(
            height,
            width,
            transparent_bg_config.chessboard_square_size,
            transparent_bg_config.chessboard_square_on_top_left,
            transparent_bg_config.chessboard_white_value,
            transparent_bg_config.chessboard_gray_value,
        )
    else:
        raise ValueError(f"Invalid background pattern: {bg_pattern}")

    alpha = img[:, :, 3]
    img_rgb = img[:, :, :3]
    alpha_normalized = alpha.astype(np.float32) / 255.0
    alpha_3d = np.stack([alpha_normalized] * 3, axis=2)
    result = alpha_3d * img_rgb + (1 - alpha_3d) * bg
    result = result.astype(np.uint8)
    return Image.fromarray(result)

def image_to_np(
    image: Image.Image,
    resize_to: tuple[int, int] | None = None,
    mode: str = "resize",
    raise_error_for_ill_resize: bool = True,
    transparent_bg_config: TransparentBgConfig | None = None,
    transparent_bg_fill_stage: Literal[
        "before_resize", "after_resize"
    ] = "before_resize",
) -> np.ndarray:
    """Convert an image to a numpy array.

    Args:
        content: The image to convert.
        resize_to: The size to resize the image to.
        mode: The mode to resize the image to.
        raise_error_for_ill_resize: Whether to raise an error for ill-sized resize.
        transparent_bg_config: The config of the transparent background. Only
            used when ``transparent_bg_fill_stage == "after_resize"`` (the
            caller is responsible for filling before resize otherwise).
        transparent_bg_fill_stage: When to composite the transparent
            background — before or after the resize step.

    Returns:
        A numpy array.
    """
    assert isinstance(image, Image.Image), "image must be a PIL Image"
    if resize_to is not None:
        if mode == "resize":
            image = image.resize(resize_to, resample=Image.Resampling.BICUBIC)
            if transparent_bg_fill_stage == "after_resize":
                image = fill_transparent_bg_with(image, transparent_bg_config)

        elif mode == "rescale_and_pad_to_center":
            scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0)
            new_width = round(image.width * scale)
            new_height = round(image.height * scale)
            if new_width == 0 or new_height == 0:
                if raise_error_for_ill_resize:
                    raise ValueError(
                        f"Invalid resize to: {resize_to}, from image size: {image.size}"
                    )
                else:
                    return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8)

            image = image.resize(
                (new_width, new_height), resample=Image.Resampling.BICUBIC
            )
            if transparent_bg_fill_stage == "after_resize":
                image = fill_transparent_bg_with(image, transparent_bg_config)
            padding_left = (resize_to[0] - new_width) // 2
            padding_right = resize_to[0] - new_width - padding_left
            padding_top = (resize_to[1] - new_height) // 2
            padding_bottom = resize_to[1] - new_height - padding_top
            image = np.asarray(image)
            image = np.pad(
                image,
                ((padding_top, padding_bottom), (padding_left, padding_right), (0, 0)),
                mode="constant",
                constant_values=0,
            )
            assert image.shape == (resize_to[1], resize_to[0], 3)

        elif mode == "rescale_and_pad_to_rightbottom":
            scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0)
            new_width = round(image.width * scale)
            new_height = round(image.height * scale)
            if new_width == 0 or new_height == 0:
                if raise_error_for_ill_resize:
                    raise ValueError(
                        f"Invalid resize to: {resize_to}, from image size: {image.size}"
                    )
                else:
                    return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8)

            image = image.resize(
                (new_width, new_height), resample=Image.Resampling.BICUBIC
            )
            if transparent_bg_fill_stage == "after_resize":
                image = fill_transparent_bg_with(image, transparent_bg_config)
            padding_right = resize_to[0] - new_width
            padding_bottom = resize_to[1] - new_height
            image = np.asarray(image)
            image = np.pad(
                image,
                ((0, padding_bottom), (0, padding_right), (0, 0)),
                mode="constant",
                constant_values=0,
            )
            assert image.shape == (resize_to[1], resize_to[0], 3)

        else:
            raise ValueError(f"Invalid mode: {mode}")

    if isinstance(image, Image.Image):
        return np.asarray(image)
    else:
        return image

def navit_patchify(pixel_values: np.ndarray, patch_size: int) -> dict[str, np.ndarray]:
    """Reshape the pixel values to a navit shape.

    Args:
        pixel_values: np.ndarray, shape (t, h, w, c)
        patch_size: int

    Returns:
        dict[str, np.ndarray]
        - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size)
        - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size)
    """
    T, H, W, C = pixel_values.shape
    assert C == 3, "pixel_values must have 3 channels"

    patches = pixel_values.reshape(
        T, H // patch_size, patch_size, W // patch_size, patch_size, C
    )
    # (T, H//patch_size, W//patch_size, C, patch_size, patch_size)
    patches = patches.transpose(0, 1, 3, 5, 2, 4)
    patches = patches.reshape(-1, C, patch_size, patch_size)
    grid_thw = np.array([T, H // patch_size, W // patch_size])
    return {"pixel_values": patches, "grid_thw": grid_thw}

def normalize(
    x: np.ndarray, mean, std_inv, pixels_dtype: np.dtype = np.float32
) -> np.ndarray:
    """Normalize the image.

    Args:
        x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255].
        mean: The mean of the image.
        std_inv: The inverse of the std of the image.
        pixels_dtype: The dtype of the image.
    Returns:
        The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype.
    """
    x = (x / 255.0).astype(pixels_dtype)
    x -= mean
    x *= std_inv
    return x