Spaces:
Running on Zero
Running on Zero
| """Direct question-Q / image-token-K pointing on Qwen's merged visual grid. | |
| The candidates are LLM image-token positions, not raw ViT keys (a different | |
| projection space). Geometry follows image_grid_thw after spatial merging, in | |
| row-major order, with no crop, label-dependent input, or learned coordinate head. | |
| """ | |
| import math | |
| import numpy as np | |
| import torch | |
| from torch.nn import functional as F | |
| from .point_probability_map import component_points | |
| from .roi_replay import REPLAY_VARIANTS | |
| VISUAL_TOKENS = "visual_tokens" | |
| ORIGINAL_IMAGE = "original_image_nine_way" | |
| def validate_pointing_config(config): | |
| variant = config.get("point_variant", "crop_refill") | |
| if variant not in ("crop_refill", VISUAL_TOKENS, ORIGINAL_IMAGE, *REPLAY_VARIANTS): | |
| raise ValueError(f"Unsupported KEV pointing method: {variant}") | |
| if variant == ORIGINAL_IMAGE and config.get("point_depth", 3) != 3: | |
| raise ValueError("Original-image nine-way pointing requires three decisions") | |
| if variant in REPLAY_VARIANTS and ( | |
| config.get('point_depth', 3) != 3 or config.get('readout_kind') != 'native_last_qk_mean' | |
| or config.get('readout_layer_indices') is not None or config.get('readout_layer_index', 31) != 31 | |
| or config.get('shared_prefix', False) or config.get('gradient_checkpointing', False)): | |
| raise ValueError('ROI replay fixes L32 mean, three decisions and no persistent cache/checkpointing') | |
| if variant != VISUAL_TOKENS: | |
| return | |
| if config.get("point_depth", 1) != 1: | |
| raise ValueError("Visual-token pointing is a single decision, not a crop tree") | |
| if config.get("shared_prefix", False): | |
| raise ValueError("Visual-token keys require the full causal row; shared_prefix is unsupported") | |
| if config.get("point_loss", "uniform_cell_ce") != "uniform_cell_ce": | |
| raise ValueError("Visual-token pointing requires uniform_cell_ce") | |
| if config.get("point_decoder", "argmax") not in ("argmax", "components"): | |
| raise ValueError("Unknown visual-token point decoder") | |
| if not 0 < config.get("point_component_threshold", 0.1) <= 1: | |
| raise ValueError("Invalid visual-token component threshold") | |
| def image_token_candidates(ids, image_token_id, grid_thw, merge_size): | |
| """One original image only; reject video, missing tokens and ambiguous grids.""" | |
| t, h, w = [int(v) for v in grid_thw] | |
| if merge_size < 1 or t != 1 or min(h, w) < 1 or h % merge_size or w % merge_size: | |
| raise ValueError("Expected one spatially mergeable still-image grid") | |
| shape = (h // merge_size, w // merge_size) | |
| positions = [i for i, token in enumerate(ids) if token == image_token_id] | |
| if len(positions) != math.prod(shape) or positions != list(range(positions[0], positions[-1] + 1)): | |
| raise ValueError("Image-token positions do not match the merged visual grid") | |
| return positions, shape | |
| def target_cells(points, shape): | |
| """Map normalized training labels to occupied cells, without modifying inputs.""" | |
| if not isinstance(points, (list, tuple)) or not points: | |
| raise ValueError("Point supervision must contain at least one normalized point") | |
| height, width = shape | |
| if min(height, width) < 1: | |
| raise ValueError("Empty visual grid") | |
| cells = set() | |
| for point in points: | |
| if (not isinstance(point, (list, tuple)) or len(point) != 2 | |
| or any(isinstance(v, bool) or not isinstance(v, (int, float)) | |
| or not math.isfinite(v) or not 0 <= v <= 1 for v in point)): | |
| raise ValueError("Point supervision must use finite normalized coordinates") | |
| x, y = point | |
| col, row = min(int(x * width), width - 1), min(int(y * height), height - 1) | |
| cells.add(row * width + col) | |
| return sorted(cells) | |
| def decision_losses(logits, packed): | |
| """Ordinary CE unchanged; multiple occupied visual cells share one QA's mass. | |
| Sparse targets avoid 0 * -inf on padded candidate slots. Coordinates in the | |
| same visual cell are intentionally indistinguishable at this resolution. | |
| """ | |
| targets = packed["targets"].to(logits.device) | |
| if "visual_target_owner" not in packed: | |
| return F.cross_entropy(logits, targets, reduction="none") | |
| log_probs = F.log_softmax(logits.float(), dim=-1) | |
| losses = -log_probs.gather(1, targets[:, None]).squeeze(1) | |
| owners = packed["visual_target_owner"].to(logits.device) | |
| slots = packed["visual_target_slot"].to(logits.device) | |
| mass = packed["visual_target_mass"].to(logits.device) | |
| visual_losses = torch.zeros_like(losses).index_add(0, owners, -log_probs[owners, slots] * mass) | |
| selected = owners.unique() | |
| return losses.index_copy(0, selected, visual_losses[selected]) | |
| def decode_visual_points(logits, shape, decoder="argmax", relative_threshold=0.1): | |
| """Decode only model scores and known geometry; never receive GT/count/masks.""" | |
| values = np.asarray(logits, dtype=np.float64) | |
| height, width = shape | |
| if min(height, width) < 1 or values.shape != (height * width,) or not np.isfinite(values).all(): | |
| raise ValueError("Expected one finite score per merged visual token") | |
| probabilities = np.exp(values - values.max()) | |
| probabilities /= probabilities.sum() | |
| grid = probabilities.reshape(height, width) | |
| if decoder == "argmax": | |
| index = int(values.argmax()) | |
| y, x = divmod(index, width) | |
| points = [[(x + 0.5) / width, (y + 0.5) / height]] | |
| elif decoder == "components": | |
| components = component_points(grid, relative_threshold=relative_threshold, connectivity=8)["components"] | |
| components.sort(key=lambda c: (-c["peak_probability"], *c["peak_cell"])) | |
| points = [[(c["peak_cell"][1] + 0.5) / width, (c["peak_cell"][0] + 0.5) / height] | |
| for c in components] | |
| else: | |
| raise ValueError(f"Unknown visual-token point decoder: {decoder}") | |
| return {"prediction": points[0], "points": points, "grid_hw": list(shape), | |
| "probabilities": grid.tolist(), "pointing_method": VISUAL_TOKENS, | |
| "point_decoder": decoder, "point_component_threshold": relative_threshold if decoder == "components" else None} | |