Spaces:
Sleeping
Sleeping
| import math | |
| import random | |
| from typing import Optional, Tuple, List, Dict, Any | |
| import os | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import torchvision | |
| from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize | |
| import onnxruntime | |
| # 如果你用 transformers 的 AutoTokenizer（推荐） | |
| from transformers import AutoTokenizer | |
| # --------------------------- | |
| # Config / 默认参数 | |
| # --------------------------- | |
| DEFAULT_IMAGE_SIZE = 224 | |
| DEFAULT_MAX_ROWS = 2 | |
| DEFAULT_MAX_COLS = 2 | |
| MIN_BLOCK_SIZE = 4 | |
| # --------------------------- | |
| # Tokenizer / Preprocess | |
| # --------------------------- | |
| def load_tokenizer(tokenizer_path: str): | |
| """加载 tokenizer（AutoTokenizer）""" | |
| tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) | |
| if tokenizer.chat_template is None: | |
| # 新版本transformer可自动加载，训练环境版本：4.51.3支持 | |
| with open(os.path.join(tokenizer_path, "chat_template.jinja"), "r") as f: | |
| tokenizer.chat_template = f.read() | |
| return tokenizer | |
| def build_image_preprocess(image_size: int = DEFAULT_IMAGE_SIZE): | |
| """返回 torchvision.transforms.Compose 的预处理 callable""" | |
| return Compose([ | |
| Resize(size=image_size, interpolation=torchvision.transforms.InterpolationMode.BICUBIC, max_size=None, antialias=True), | |
| CenterCrop(size=(image_size, image_size)), | |
| lambda img: img.convert("RGB"), | |
| ToTensor(), | |
| Normalize(mean=(0.48145466, 0.4578275, 0.40821073), | |
| std=(0.26862954, 0.26130258, 0.27577711)) | |
| ]) | |
| # --------------------------- | |
| # RoPE频率预计算（precompute_freqs_cis） | |
| # --------------------------- | |
| def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, | |
| rope_scaling: Optional[dict] = None) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| 计算 RoPE 的 cos 和 sin 表 | |
| 返回： | |
| freqs_cos: (end, dim) | |
| freqs_sin: (end, dim) | |
| """ | |
| freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) | |
| if rope_scaling is not None: | |
| orig_max = rope_scaling.get("original_max_position_embeddings", 2048) | |
| factor = rope_scaling.get("factor", 4) | |
| beta_fast = rope_scaling.get("beta_fast", 4.0) | |
| beta_slow = rope_scaling.get("beta_slow", 1.0) | |
| if end / orig_max > 1.0: | |
| corr_dim = next((i for i in range(dim // 2) if 2 * math.pi / freqs[i] > orig_max), dim // 2) | |
| power = torch.arange(0, dim // 2, device=freqs.device).float() / max(dim // 2 - 1, 1) | |
| beta = beta_slow + (beta_fast - beta_slow) * power | |
| scale = torch.where(torch.arange(dim // 2, device=freqs.device) < corr_dim, | |
| (beta * factor - beta + 1) / (beta * factor), | |
| 1.0 / factor) | |
| freqs = freqs * scale | |
| t = torch.arange(end, device=freqs.device) | |
| freqs = torch.outer(t, freqs).float() # (end, dim/2) | |
| freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) | |
| freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) | |
| return freqs_cos, freqs_sin | |
| # --------------------------- | |
| # 图像自适应切分（adaptive_square_split） | |
| # --------------------------- | |
| def calculate_optimal_split_with_fixed_max(width: int, height: int, max_rows: int, max_cols: int) -> Tuple[int, int, int]: | |
| """ | |
| 计算最佳切分（返回 rows, cols, block_size） | |
| block_size 会向下取整到 16 的倍数，最小为 MIN_BLOCK_SIZE。 | |
| """ | |
| best_rows = 1 | |
| best_cols = 1 | |
| best_block_size = 0 | |
| best_coverage = 0.0 | |
| # 方案1: 固定行数为 max_rows，自适应列数 | |
| rows_fixed = max_rows | |
| for cols in range(1, max_cols + 1): | |
| block_width = width // cols | |
| block_height = height // rows_fixed | |
| square_size = min(block_width, block_height) | |
| if square_size <= 0: | |
| continue | |
| coverage = (cols * square_size) * (rows_fixed * square_size) / (width * height) | |
| if coverage > best_coverage or (coverage == best_coverage and square_size > best_block_size): | |
| best_rows, best_cols, best_block_size, best_coverage = rows_fixed, cols, square_size, coverage | |
| # 方案2: 固定列数为 max_cols，自适应行数 | |
| cols_fixed = max_cols | |
| for rows in range(1, max_rows + 1): | |
| block_width = width // cols_fixed | |
| block_height = height // rows | |
| square_size = min(block_width, block_height) | |
| if square_size <= 0: | |
| continue | |
| coverage = (cols_fixed * square_size) * (rows * square_size) / (width * height) | |
| if coverage > best_coverage or (coverage == best_coverage and square_size > best_block_size): | |
| best_rows, best_cols, best_block_size, best_coverage = rows, cols_fixed, square_size, coverage | |
| # 方案3: 两者都达到最大 | |
| block_width = width // max_cols | |
| block_height = height // max_rows | |
| square_size = min(block_width, block_height) | |
| if square_size > 0: | |
| coverage = (max_cols * square_size) * (max_rows * square_size) / (width * height) | |
| if coverage > best_coverage or (coverage == best_coverage and square_size > best_block_size): | |
| best_rows, best_cols, best_block_size, best_coverage = max_rows, max_cols, square_size, coverage | |
| # 对齐到 16 的倍数并保证最小值 | |
| best_block_size = max(MIN_BLOCK_SIZE, (best_block_size // 16) * 16) | |
| return best_rows, best_cols, best_block_size | |
| def adaptive_square_split(image: Image.Image, max_rows: int = DEFAULT_MAX_ROWS, max_cols: int = DEFAULT_MAX_COLS | |
| ) -> Tuple[List[Image.Image], int, int, int]: | |
| """ | |
| 将 PIL Image 自适应切分为正方形块，返回 (blocks_list, rows, cols, block_size) | |
| blocks_list 是按行主序的块列表（可能少于 max_rows*max_cols） | |
| """ | |
| width, height = image.size | |
| rows, cols, block_size = calculate_optimal_split_with_fixed_max(width, height, max_rows, max_cols) | |
| blocks = [] | |
| for r in range(rows): | |
| for c in range(cols): | |
| left = c * block_size | |
| upper = r * block_size | |
| right = left + block_size | |
| lower = upper + block_size | |
| blocks.append(image.crop((left, upper, right, lower))) | |
| return blocks, rows, cols, block_size | |
| # --------------------------- | |
| # 特殊 token 准备 | |
| # --------------------------- | |
| def prepare_special_tokens(tokenizer, max_rows: int = DEFAULT_MAX_ROWS, max_cols: int = DEFAULT_MAX_COLS) -> Dict[str, int]: | |
| """ | |
| 返回特殊 token id 的 dict，包含 <global-img>, <fake_token_around_image>, <image>, 以及 <row_i_col_j> | |
| """ | |
| special = { | |
| "<global-img>": tokenizer.convert_tokens_to_ids("<global-img>"), | |
| "<fake_token_around_image>": tokenizer.convert_tokens_to_ids("<fake_token_around_image>"), | |
| "<image>": tokenizer.convert_tokens_to_ids("<image>"), | |
| } | |
| for i in range(max_rows): | |
| for j in range(max_cols): | |
| special[f"<row_{i + 1}_col_{j + 1}>"] = tokenizer.convert_tokens_to_ids(f"<row_{i + 1}_col_{j + 1}>") | |
| return special | |
| # --------------------------- | |
| # 将图像切块、填充、stack 为模型输入张量 | |
| # --------------------------- | |
| def prepare_image_patches(image: Image.Image, preprocess_fn, max_rows: int = DEFAULT_MAX_ROWS, max_cols: int = DEFAULT_MAX_COLS | |
| ) -> Tuple[torch.Tensor, List[int]]: | |
| """ | |
| image: PIL.Image | |
| preprocess_fn: callable that maps PIL.Image -> tensor(C,H,W) | |
| 返回： | |
| pixel_values: torch.Tensor, shape (num_patches + 1, C, H, W) -- 最后一个是 full image 原图 | |
| mask_token_ids: list[int] -- 当某个位置为空时，对应的 row_col token id 列表（未去重） | |
| """ | |
| blocks, rows, cols, block_size = adaptive_square_split(image, max_rows=max_rows, max_cols=max_cols) | |
| patch_num = len(blocks) | |
| pad_num = max_rows * max_cols - patch_num | |
| mask_token_id_list = [] | |
| patch_tensors = [] | |
| if pad_num > 0: | |
| # 以行主序填充: 若某个位置超出 rows 或 cols，则用零张量并记录对应的 row_col token id（由调用者映射） | |
| for i in range(max_rows): | |
| for j in range(max_cols): | |
| if i >= rows or j >= cols: | |
| patch_tensors.append(torch.zeros_like(preprocess_fn(image))) | |
| # mask token id 由调用者生成/映射，这里只记录一个占位（具体 id 值需外部映射） | |
| # 返回时，调用者会在文本中找到对应的 special token 的位置并进行 attention mask 操作 | |
| mask_token_id_list.append((i, j)) | |
| else: | |
| patch_tensors.append(preprocess_fn(blocks[i * cols + j])) | |
| else: | |
| patch_tensors = [preprocess_fn(b) for b in blocks] | |
| # 最后附加 full image 的 pixel_values（和你原来逻辑一致） | |
| full_image_tensor = preprocess_fn(image) | |
| pixel_values = torch.stack(patch_tensors + [full_image_tensor], dim=0) # (N_patches+1, C, H, W) | |
| return pixel_values, mask_token_id_list | |
| # --------------------------- | |
| # 在 token 流中构建 image placeholder（原始的占位 token 串） | |
| # --------------------------- | |
| def construct_image_placeholders(special_tokens: Dict[str, int], max_rows: int = DEFAULT_MAX_ROWS, max_cols: int = DEFAULT_MAX_COLS, | |
| n_image_tokens_per_patch: int = 49) -> str: | |
| """ | |
| 生成一个示例占位字符串，便于拼接到 prompt 中。 | |
| 返回一个包含多个占位符的字符串 (str)。 | |
| """ | |
| image_place_holder = random.choice(["图片如下：", "如下所示的图片:", "请见下面这张图:", "如下图显示:", "参考下方图片:", "图示如下:"]) | |
| for row in range(max_rows): | |
| for col in range(max_cols): | |
| image_place_holder += f"<fake_token_around_image><row_{row + 1}_col_{col + 1}>" | |
| image_place_holder += "<image>" * n_image_tokens_per_patch | |
| # 全局图像块（最后） | |
| image_place_holder += f"<fake_token_around_image><global-img>{'<image>' * n_image_tokens_per_patch}<fake_token_around_image>" | |
| return image_place_holder | |
| # --------------------------- | |
| # 寻找 token 序列中 image 标记出现的位置（用于 attention mask 修改） | |
| # --------------------------- | |
| def find_indices(tokens: torch.Tensor) -> Optional[Dict[int, Dict[int, List[Tuple[int, int]]]]]: | |
| """ | |
| 输入 tokens: shape (B, T) 的 tensor | |
| 返回结构: | |
| results = { batch_index: { k: [(start_idx, end_idx), ...], ... }, ... } | |
| 其中 k 对应 image token 的索引（函数里预设 image_id 列表），返回的 start_idx/end_idx 为占位段在 tokens 中的 start/end（包含） | |
| 说明：此方法沿用了你原来的匹配模式（匹配 [<fake>, <row_i_col_j>] 以及 [<fake>, <global-img>]） | |
| """ | |
| B, T = tokens.size() | |
| # 这里使用与原代码一致的 id 序列（如果 tokenizer 中不同，请改这里） | |
| # image_ids = [[3, i] for i in range(6, 22)] + [[3, 4]] # 预设 pattern | |
| image_ids = [[3, 6], [3, 7], [3, 10], [3, 11],[3, 4]] | |
| image_ids_tensor = torch.tensor(image_ids, device=tokens.device) | |
| len_image_ids = image_ids_tensor.size(1) | |
| if len_image_ids > tokens.size(1): | |
| return None | |
| tokens_view = tokens.unfold(1, len_image_ids, 1) # (B, T - len_image_ids +1, len_image_ids) | |
| matches = [] | |
| for image_id_tensor in image_ids_tensor: | |
| match = (tokens_view == image_id_tensor).all(dim=2) # (B, T-len+1) | |
| matches.append(match) | |
| results = {} | |
| for b in range(B): | |
| batch_res = {} | |
| for k, m in enumerate(matches): | |
| idxs = m[b].nonzero(as_tuple=True)[0] | |
| if len(idxs) > 0: | |
| batch_res[k] = [(i.item() + 2, i.item() + 50) for i in idxs] | |
| if batch_res: | |
| results[b] = batch_res | |
| return results or None | |
| # --------------------------- | |
| # ONNX Session helpers | |
| # --------------------------- | |
| def create_onnx_session(path: str, intra_threads: int = 1) -> onnxruntime.InferenceSession: | |
| opts = onnxruntime.SessionOptions() | |
| opts.intra_op_num_threads = intra_threads | |
| opts.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL | |
| return onnxruntime.InferenceSession(path, sess_options=opts) | |
| # --------------------------- | |
| # Prefill 阶段（将视觉嵌入插入并运行一次 LLM） | |
| # --------------------------- | |
| def prefill_llm(vision_session: onnxruntime.InferenceSession, | |
| embed_tokens_session: onnxruntime.InferenceSession, | |
| llm_session: onnxruntime.InferenceSession, | |
| pixel_values: torch.Tensor, | |
| input_ids: torch.Tensor, | |
| attention_mask: torch.Tensor, | |
| freqs_cos: torch.Tensor, | |
| freqs_sin: torch.Tensor, | |
| special_tokens: Dict[str, int], | |
| seqlen: int, | |
| device: str = "cpu") -> Dict[str, Any]: | |
| """ | |
| 完成 prefill 步骤： | |
| 1) 通过 vision_session 获得视觉嵌入 deepstack_embeds | |
| 2) 通过 embed_tokens_session 获得 token embedding（或直接使用输入 hidden states） | |
| 3) 将视觉嵌入插入 hidden stream（替换占位 token 段） | |
| 4) 调用 llm_session.run 一次，得到 logits、hidden_states、present_keys、present_values | |
| 返回 dict: | |
| { | |
| "logits": np.ndarray, | |
| "hidden_states": np.ndarray, | |
| "present_keys": np.ndarray, | |
| "present_values": np.ndarray | |
| } | |
| """ | |
| # 1) vision embed | |
| ort_inputs_vis = {"inputs": pixel_values.numpy()} | |
| deepstack_embeds = vision_session.run(["deepstack_embeds"], ort_inputs_vis)[0] # e.g. (B, P, L_patch, D) | |
| # 2) token embed | |
| ort_inputs_emb = {"input_ids": input_ids.numpy()} | |
| embed_tokens = embed_tokens_session.run(["embed_tokens"], ort_inputs_emb)[0] # (B, T, D) | |
| # 3) 找到 image placeholder 在 token 中的位置并替换 | |
| image_batch_indices = find_indices(input_ids) | |
| B = input_ids.shape[0] | |
| seqlen = seqlen | |
| new_h = [] | |
| for i in range(B): | |
| h_i = embed_tokens[i] # np array (T, D) | |
| image_indices = image_batch_indices.get(i, {}) if image_batch_indices else {} | |
| # image_indices: {k: [(start,end), ...], ...} | |
| # deepstack_embeds: assume shape (B, P, L_patch, D), P = number_of_image_patches + global | |
| for tki, index_list in image_indices.items(): | |
| # tki 对应 deepstack_embeds 第二维索引 | |
| vision_proj_i = deepstack_embeds[i][tki] # (L_patch, D) | |
| # 取第一个匹配段 | |
| start_idx, end_idx = index_list[0] | |
| # 将 h_i 中 start_idx..end_idx 替换为 vision_proj_i，并截断到 seqlen | |
| # 注意这里我们使用 numpy concat（h_i 是 numpy） | |
| h_i = np.concatenate((h_i[:start_idx], vision_proj_i, h_i[end_idx + 1:]), axis=0)[:seqlen] | |
| new_h.append(h_i) | |
| hidden_states = np.stack(new_h, axis=0) # (B, seqlen, D) | |
| # 4) 呼叫 llm.onnx 做一次前向（prefill） | |
| # past_keys/past_values 用空的 shape（按模型要求） | |
| # 这里 past keys/values 的 shape 需与模型期望一致，示例用随机的 0 长度数组作为占位 | |
| # 如模型要求具体形状，请在调用方准备 | |
| past_keys = np.zeros([8, 0, 2, 64], dtype=np.float32) | |
| past_values = np.zeros([8, 0, 2, 64], dtype=np.float32) | |
| cos_pe = freqs_cos[0: seqlen].numpy() | |
| sin_pe = freqs_sin[0: seqlen].numpy() | |
| ort_inputs_llm = { | |
| "input_ids": hidden_states.astype(np.float32), | |
| "attention_mask": attention_mask.numpy(), | |
| "cos_pe": cos_pe.astype(np.float32), | |
| "sin_pe": sin_pe.astype(np.float32), | |
| "past_keys": past_keys, | |
| "past_values": past_values | |
| } | |
| logits, hidden_states_out, present_keys, present_values = llm_session.run( | |
| ["logits", "hidden_states", "present_keys", "present_values"], ort_inputs_llm | |
| ) | |
| return { | |
| "logits": logits, | |
| "hidden_states": hidden_states_out, | |
| "present_keys": present_keys, | |
| "present_values": present_values | |
| } | |
| # --------------------------- | |
| # Next-token 自回归生成（基于 present keys/values） | |
| # --------------------------- | |
| def generate_autoregressive(llm_session: onnxruntime.InferenceSession, | |
| embed_tokens_session: onnxruntime.InferenceSession, | |
| tokenizer, | |
| initial_present: Dict[str, np.ndarray], | |
| start_token_id: int, | |
| freqs_cos: torch.Tensor, | |
| freqs_sin: torch.Tensor, | |
| attention_mask: np.ndarray, | |
| max_new_tokens: int = 128, | |
| eos_token_id: int = 2, | |
| start_pos: int = None, | |
| output_queue=None): | |
| """ | |
| 基于 prefill 返回的 present_keys/present_values 进行自回归生成。 | |
| 每一步： | |
| - 用 embed_tokens_session 获取新 token 的 embedding | |
| - 用 llm_session 传入 present keys/values 并得到新的 present keys/values 与 logits | |
| - 选取最大 logit（argmax）作为下一个 token（你可替换为 sampling 策略） | |
| 如果提供了 output_queue，则将每个生成的文本片段放入队列。 | |
| 注意：present keys/values 的名称与 shape 与模型实现相关，确保和模型一致。 | |
| """ | |
| present_keys = initial_present["present_keys"] | |
| present_values = initial_present["present_values"] | |
| present_keys = present_keys | |
| present_values = present_values | |
| token_id = int(start_token_id) | |
| if start_pos is None: | |
| start_pos = attention_mask.shape[1] | |
| generated_ids = [] | |
| for step in range(max_new_tokens): | |
| # 解码当前 token | |
| decoded = tokenizer.decode(token_id, skip_special_tokens=True, clean_up_tokenization_spaces=True) | |
| # 如果提供了输出队列，放入生成的文本 | |
| if output_queue is not None: | |
| output_queue.put(decoded) | |
| # 更新 attention mask | |
| attention_mask = np.concatenate([attention_mask, np.array([[1]], dtype=np.int64)], axis=1) | |
| # embed 当前 token | |
| embed_tokens = embed_tokens_session.run(["embed_tokens"], {"input_ids": np.array([[token_id]], dtype=np.int64)})[0] | |
| cos_pe = freqs_cos[start_pos: start_pos + 1].numpy() | |
| sin_pe = freqs_sin[start_pos: start_pos + 1].numpy() | |
| ort_inputs = { | |
| "input_ids": embed_tokens.astype(np.float32), | |
| "attention_mask": attention_mask, | |
| "cos_pe": cos_pe.astype(np.float32), | |
| "sin_pe": sin_pe.astype(np.float32), | |
| "past_keys": present_keys, | |
| "past_values": present_values | |
| } | |
| logits, hidden_states, present_keys, present_values = llm_session.run( | |
| ["logits", "hidden_states", "present_keys", "present_values"], ort_inputs | |
| ) | |
| token_id = int(np.argmax(logits[:, -1, :], axis=-1)[0]) | |
| generated_ids.append(token_id) | |
| if token_id == eos_token_id: | |
| break | |
| start_pos += 1 | |
| # 发送生成结束信号 | |
| if output_queue is not None: | |
| output_queue.put(None) | |
| return generated_ids | |
| def main_example(): | |
| tokenizer_path = "./custom_tokenizer" | |
| tokenizer = load_tokenizer(tokenizer_path) | |
| preprocess = build_image_preprocess(DEFAULT_IMAGE_SIZE) | |
| # image | |
| image_path = "/Users/hulk/Downloads/coco128/images/train2017/000000000165.jpg" | |
| image = Image.open(image_path).convert("RGB") | |
| # special tokens | |
| special_tokens = prepare_special_tokens(tokenizer, max_rows=4, max_cols=4) | |
| pixel_values, mask_positions = prepare_image_patches(image, preprocess, max_rows=4, max_cols=4) | |
| # 构造 prompt + image placeholders（假设 tokenizer 支持 apply_chat_template） | |
| query = "图片中的人在做什么。" | |
| messages = [ | |
| {"role": "system", "content": "你是一个多模态AI助手，能够理解图片和文本信息."}, | |
| {"role": "user", "content": query + construct_image_placeholders(special_tokens)} | |
| ] | |
| inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True) | |
| input_ids = inputs["input_ids"] | |
| attention_mask = inputs["attention_mask"] | |
| # precompute RoPE | |
| freqs_cos, freqs_sin = precompute_freqs_cis(dim=64, end=32768, rope_base=1e6) | |
| # create onnx sessions | |
| vision_session = create_onnx_session("./onnx_model/vision_encoder.onnx", intra_threads=2) | |
| embed_tokens_session = create_onnx_session("./onnx_model/embed_tokens.onnx", intra_threads=2) | |
| llm_session = create_onnx_session("./onnx_model/llm.onnx", intra_threads=2) | |
| # prefill | |
| seqlen = input_ids.shape[1] | |
| prefill_out = prefill_llm( | |
| vision_session=vision_session, | |
| embed_tokens_session=embed_tokens_session, | |
| llm_session=llm_session, | |
| pixel_values=pixel_values, | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| freqs_cos=freqs_cos, | |
| freqs_sin=freqs_sin, | |
| special_tokens=special_tokens, | |
| seqlen=seqlen | |
| ) | |
| # start token id = argmax last logit | |
| start_token_id = int(np.argmax(prefill_out["logits"][:, -1, :], axis=-1)[0]) | |
| generated = generate_autoregressive( | |
| llm_session=llm_session, | |
| embed_tokens_session=embed_tokens_session, | |
| tokenizer=tokenizer, | |
| initial_present={"present_keys": prefill_out["present_keys"], "present_values": prefill_out["present_values"]}, | |
| start_token_id=start_token_id, | |
| freqs_cos=freqs_cos, | |
| freqs_sin=freqs_sin, | |
| attention_mask=attention_mask.numpy(), | |
| max_new_tokens=128, | |
| eos_token_id=2, | |
| start_pos=seqlen | |
| ) | |
| if __name__ == "__main__": | |
| # main_example() | |
| pass |