from __future__ import annotations from typing import Dict, Optional, Tuple import torch import torch.nn.functional as F from .model import _color_tone, _grayscale, _sobel_edges, ArtistStyleModel def _normalize_map(x: torch.Tensor) -> torch.Tensor: x = x - x.amin(dim=(-2, -1), keepdim=True) x = x / x.amax(dim=(-2, -1), keepdim=True).clamp_min(1e-6) return x def _dinov3_cls_attention_map(backbone, raw_view: torch.Tensor) -> torch.Tensor: normalized = ( raw_view - backbone.pixel_mean.to(dtype=raw_view.dtype, device=raw_view.device) ) / backbone.pixel_std.to(dtype=raw_view.dtype, device=raw_view.device) dino = backbone.backbone tokens, patch_hw = dino.prepare_tokens_with_masks(normalized) for block_idx, block in enumerate(dino.blocks): rope = dino.rope_embed(H=patch_hw[0], W=patch_hw[1]) if dino.rope_embed is not None else None if block_idx == len(dino.blocks) - 1: attn_input = block.norm1(tokens) attn_module = block.attn qkv = attn_module.qkv(attn_input) batch, token_count, _ = qkv.shape channel_count = attn_module.qkv.in_features qkv = qkv.reshape( batch, token_count, 3, attn_module.num_heads, channel_count // attn_module.num_heads, ) q, k, _ = torch.unbind(qkv, 2) q, k = [tensor.transpose(1, 2) for tensor in (q, k)] if rope is not None: q, k = attn_module.apply_rope(q, k, rope) logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * float(attn_module.scale) attention = torch.softmax(logits, dim=-1) patch_start = 1 + int(getattr(dino, "n_storage_tokens", 0) or 0) cls_to_patch = attention[:, :, 0, patch_start:].mean(dim=1) heatmap = cls_to_patch.reshape(batch, 1, patch_hw[0], patch_hw[1]) heatmap = F.interpolate(heatmap, size=raw_view.shape[-2:], mode="bilinear", align_corners=False) return _normalize_map(heatmap.to(dtype=raw_view.dtype)) tokens = block(tokens, rope) return raw_view.new_zeros((raw_view.size(0), 1, raw_view.size(-2), raw_view.size(-1))) def _apply_attention_gate(heatmap: torch.Tensor, attention_map: torch.Tensor) -> torch.Tensor: attention_gate = (0.2 + attention_map.float()).pow(0.75).to(dtype=heatmap.dtype) return _normalize_map(heatmap * attention_gate) def _crop_attention_map( attention_map: torch.Tensor, normalized_box: Tuple[float, float, float, float], output_size: tuple[int, int], ) -> torch.Tensor: _, _, height, width = attention_map.shape x1, y1, x2, y2 = normalized_box left = max(0, min(width - 1, int(round(x1 * width)))) top = max(0, min(height - 1, int(round(y1 * height)))) right = max(left + 1, min(width, int(round(x2 * width)))) bottom = max(top + 1, min(height, int(round(y2 * height)))) cropped = attention_map[:, :, top:bottom, left:right] return F.interpolate(cropped, size=output_size, mode="bilinear", align_corners=False) def _renorm_weights(x: torch.Tensor, dim: int = -1) -> torch.Tensor: return x / x.sum(dim=dim, keepdim=True).clamp_min(1e-6) def _branch_patch_prior(branch_name: str, branch_module, raw_view: torch.Tensor, patch_tokens: torch.Tensor, patch_hw: tuple[int, int]) -> torch.Tensor: if patch_tokens.size(1) == 0: return raw_view.new_zeros((raw_view.size(0), 0)) if branch_name == "texture": prompt = F.normalize(branch_module.token_prompt.float(), dim=0) patch_norm = F.normalize(patch_tokens.float(), dim=-1) return torch.softmax(torch.einsum("bnd,d->bn", patch_norm, prompt), dim=1) if branch_name == "line": edge_map = _grayscale(_sobel_edges(raw_view.float())) weights = F.adaptive_avg_pool2d(edge_map, patch_hw).flatten(1) return weights / weights.sum(dim=1, keepdim=True).clamp_min(1e-6) if branch_name == "color": color_map = _color_tone(raw_view.float())[:, 1:].pow(2).sum(dim=1, keepdim=True).sqrt() weights = F.adaptive_avg_pool2d(color_map, patch_hw).flatten(1) return weights / weights.sum(dim=1, keepdim=True).clamp_min(1e-6) weights = raw_view.new_ones((raw_view.size(0), patch_tokens.size(1)), dtype=torch.float32) return weights / weights.sum(dim=1, keepdim=True).clamp_min(1e-6) def _flip_if_needed(x: torch.Tensor, enabled: bool) -> torch.Tensor: if not enabled: return x return torch.flip(x, dims=(-1,)) def _aggregate_outputs(outputs_list: list[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: if len(outputs_list) == 1: return outputs_list[0] weighted_branch_projected = torch.stack( [output["weighted_branch_projected_embeddings"] for output in outputs_list], dim=0, ).mean(dim=0) aggregated = dict(outputs_list[0]) aggregated["branch_embeddings"] = F.normalize( torch.stack([output["branch_embeddings"] for output in outputs_list], dim=0).mean(dim=0).float(), dim=-1, ) aggregated["branch_projected_embeddings"] = F.normalize( torch.stack([output["branch_projected_embeddings"] for output in outputs_list], dim=0).mean(dim=0).float(), dim=-1, ) aggregated["weighted_branch_projected_embeddings"] = weighted_branch_projected aggregated["embedding"] = F.normalize(weighted_branch_projected.flatten(1).float(), dim=-1) aggregated["stacked_view_embeddings"] = F.normalize( torch.stack([output["stacked_view_embeddings"] for output in outputs_list], dim=0).mean(dim=0).float(), dim=-1, ) aggregated["stacked_view_weights"] = _renorm_weights( torch.stack([output["stacked_view_weights"] for output in outputs_list], dim=0).mean(dim=0), dim=-1, ) aggregated["branch_weights"] = _renorm_weights( torch.stack([output["branch_weights"] for output in outputs_list], dim=0).mean(dim=0), dim=-1, ) aggregated["view_mask"] = outputs_list[0]["view_mask"] aggregated["effective_view_mask"] = ( torch.stack([output["effective_view_mask"] for output in outputs_list], dim=0).mean(dim=0) > 0 ).to(dtype=outputs_list[0]["effective_view_mask"].dtype) aggregated["branch_mask"] = outputs_list[0]["branch_mask"] aggregated["effective_branch_mask"] = ( torch.stack([output["effective_branch_mask"] for output in outputs_list], dim=0).mean(dim=0) > 0 ).to(dtype=outputs_list[0]["effective_branch_mask"].dtype) aggregated["view_weights"] = { name: aggregated["stacked_view_weights"][:, idx] for idx, name in enumerate(["structure", "texture", "line", "color"]) } return aggregated def _encode_query( model: ArtistStyleModel, full: torch.Tensor, face: torch.Tensor, eye: torch.Tensor, view_mask: torch.Tensor, use_tta: bool = False, ) -> Dict[str, torch.Tensor]: variants = [(False, False, False)] if use_tta: variants.append((True, True, True)) outputs_list = [] for flip_full, flip_face, flip_eye in variants: outputs_list.append( model( _flip_if_needed(full, flip_full), _flip_if_needed(face, flip_face), _flip_if_needed(eye, flip_eye), view_mask=view_mask, ) ) return _aggregate_outputs(outputs_list) def similarity_breakdown(query_outputs: Dict[str, torch.Tensor], reference_outputs: Dict[str, torch.Tensor], branch_names: list[str]) -> Dict[str, object]: query_chunks = query_outputs["weighted_branch_projected_embeddings"] reference_chunks = reference_outputs["weighted_branch_projected_embeddings"] query_raw = query_chunks.flatten(1) reference_raw = reference_chunks.flatten(1) denom = query_raw.norm(dim=1) * reference_raw.norm(dim=1) branch_scores = (query_chunks * reference_chunks).sum(dim=-1) / denom.unsqueeze(-1).clamp_min(1e-6) total_similarity = (query_outputs["embedding"] * reference_outputs["embedding"]).sum(dim=-1) query_branch = query_outputs["branch_embeddings"] reference_branch = reference_outputs["branch_embeddings"] query_view = query_outputs["stacked_view_embeddings"] reference_view = reference_outputs["stacked_view_embeddings"] query_view_weights = query_outputs["stacked_view_weights"] reference_view_weights = reference_outputs["stacked_view_weights"] view_scores = [] for branch_idx in range(len(branch_names)): query_target = reference_branch[:, branch_idx].unsqueeze(1) reference_target = query_branch[:, branch_idx].unsqueeze(1) query_score = (query_view[:, branch_idx] * query_target).sum(dim=-1) * query_view_weights[:, branch_idx] reference_score = (reference_view[:, branch_idx] * reference_target).sum(dim=-1) * reference_view_weights[:, branch_idx] view_scores.append(0.5 * (query_score + reference_score)) view_scores = torch.stack(view_scores, dim=1) return { "total_similarity": total_similarity, "branch_contributions": {name: branch_scores[:, idx] for idx, name in enumerate(branch_names)}, "view_contributions": {name: view_scores[:, idx] for idx, name in enumerate(branch_names)}, } def _descriptor_to_device(reference_descriptor: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: result = {} for key, value in reference_descriptor.items(): if torch.is_tensor(value): tensor = value.to(device) if tensor.dim() >= 1: tensor = tensor.unsqueeze(0) result[key] = tensor return result @torch.no_grad() def explain_against_reference( model: ArtistStyleModel, query_full: torch.Tensor, query_face: torch.Tensor, query_eye: torch.Tensor, query_view_mask: torch.Tensor, reference_descriptor: Dict[str, torch.Tensor], view_attention_boxes: Optional[Dict[str, Tuple[float, float, float, float]]] = None, use_tta: bool = False, ) -> Dict[str, object]: was_training = model.training model.eval() try: query_outputs = _encode_query( model, query_full, query_face, query_eye, view_mask=query_view_mask, use_tta=use_tta, ) reference_outputs = _descriptor_to_device(reference_descriptor, query_outputs["embedding"].device) breakdown = similarity_breakdown(query_outputs, reference_outputs, model.branch_names) query_views = [query_full, query_face, query_eye] reference_branch_embeddings = reference_outputs["branch_embeddings"] view_attention_boxes = view_attention_boxes or {} branch_heatmaps: Dict[str, Dict[str, Optional[torch.Tensor]]] = {} combined_view_heatmaps = {"full": None, "face": None, "eye": None} attention_cache: Dict[str, torch.Tensor] = {} for branch_idx, branch_name in enumerate(model.branch_names): branch_module = model.branches[branch_name] branch_heatmaps[branch_name] = {} branch_weight = breakdown["branch_contributions"][branch_name].view(-1, 1, 1, 1) for view_idx, view_name in enumerate(["full", "face", "eye"]): if query_view_mask[:, view_idx].sum() == 0: branch_heatmaps[branch_name][view_name] = None continue raw_view = query_views[view_idx] feature_pack = model.backbone(raw_view) patch_tokens = feature_pack["patch_tokens"] patch_hw = feature_pack["patch_hw"] projected_patches = F.normalize(branch_module.patch_proj(patch_tokens.float()), dim=-1) target = F.normalize(reference_branch_embeddings[:, branch_idx], dim=-1) patch_scores = torch.einsum("bnd,bd->bn", projected_patches, target) prior = _branch_patch_prior(branch_name, branch_module, raw_view, patch_tokens, patch_hw) patch_scores = patch_scores * prior heatmap = patch_scores.view(raw_view.size(0), 1, patch_hw[0], patch_hw[1]) heatmap = F.interpolate(heatmap, size=raw_view.shape[-2:], mode="bilinear", align_corners=False) heatmap = _normalize_map(heatmap) if "full" not in attention_cache: attention_cache["full"] = _dinov3_cls_attention_map(model.backbone, query_full) if view_name == "full": attention_map = attention_cache["full"] heatmap = _apply_attention_gate(heatmap, attention_map) elif view_name in view_attention_boxes: attention_map = _crop_attention_map( attention_cache["full"], view_attention_boxes[view_name], output_size=raw_view.shape[-2:], ) heatmap = _apply_attention_gate(heatmap, attention_map) branch_heatmaps[branch_name][view_name] = heatmap weighted_map = heatmap * branch_weight if combined_view_heatmaps[view_name] is None: combined_view_heatmaps[view_name] = weighted_map else: combined_view_heatmaps[view_name] = combined_view_heatmaps[view_name] + weighted_map for view_name, heatmap in combined_view_heatmaps.items(): if heatmap is not None: combined_view_heatmaps[view_name] = _normalize_map(heatmap) return { "query_outputs": query_outputs, "reference_outputs": reference_outputs, **breakdown, "branch_heatmaps": branch_heatmaps, "combined_view_heatmaps": combined_view_heatmaps, } finally: model.train(was_training) @torch.no_grad() def explain_pair( model: ArtistStyleModel, query_full: torch.Tensor, query_face: torch.Tensor, query_eye: torch.Tensor, query_view_mask: torch.Tensor, reference_full: torch.Tensor, reference_face: torch.Tensor, reference_eye: torch.Tensor, reference_view_mask: torch.Tensor, view_attention_boxes: Optional[Dict[str, Tuple[float, float, float, float]]] = None, use_tta: bool = False, ) -> Dict[str, object]: was_training = model.training model.eval() try: query_outputs = _encode_query( model, query_full, query_face, query_eye, view_mask=query_view_mask, use_tta=use_tta, ) reference_outputs = model(reference_full, reference_face, reference_eye, view_mask=reference_view_mask) breakdown = similarity_breakdown(query_outputs, reference_outputs, model.branch_names) query_views = [query_full, query_face, query_eye] reference_branch_embeddings = reference_outputs["branch_embeddings"] view_attention_boxes = view_attention_boxes or {} branch_heatmaps: Dict[str, Dict[str, Optional[torch.Tensor]]] = {} combined_view_heatmaps = {"full": None, "face": None, "eye": None} attention_cache: Dict[str, torch.Tensor] = {} for branch_idx, branch_name in enumerate(model.branch_names): branch_module = model.branches[branch_name] branch_heatmaps[branch_name] = {} branch_weight = breakdown["branch_contributions"][branch_name].view(-1, 1, 1, 1) for view_idx, view_name in enumerate(["full", "face", "eye"]): if query_view_mask[:, view_idx].sum() == 0: branch_heatmaps[branch_name][view_name] = None continue raw_view = query_views[view_idx] feature_pack = model.backbone(raw_view) patch_tokens = feature_pack["patch_tokens"] patch_hw = feature_pack["patch_hw"] projected_patches = F.normalize(branch_module.patch_proj(patch_tokens.float()), dim=-1) target = F.normalize(reference_branch_embeddings[:, branch_idx], dim=-1) patch_scores = torch.einsum("bnd,bd->bn", projected_patches, target) prior = _branch_patch_prior(branch_name, branch_module, raw_view, patch_tokens, patch_hw) patch_scores = patch_scores * prior heatmap = patch_scores.view(raw_view.size(0), 1, patch_hw[0], patch_hw[1]) heatmap = F.interpolate(heatmap, size=raw_view.shape[-2:], mode="bilinear", align_corners=False) heatmap = _normalize_map(heatmap) if "full" not in attention_cache: attention_cache["full"] = _dinov3_cls_attention_map(model.backbone, query_full) if view_name == "full": attention_map = attention_cache["full"] heatmap = _apply_attention_gate(heatmap, attention_map) elif view_name in view_attention_boxes: attention_map = _crop_attention_map( attention_cache["full"], view_attention_boxes[view_name], output_size=raw_view.shape[-2:], ) heatmap = _apply_attention_gate(heatmap, attention_map) branch_heatmaps[branch_name][view_name] = heatmap weighted_map = heatmap * branch_weight if combined_view_heatmaps[view_name] is None: combined_view_heatmaps[view_name] = weighted_map else: combined_view_heatmaps[view_name] = combined_view_heatmaps[view_name] + weighted_map for view_name, heatmap in combined_view_heatmaps.items(): if heatmap is not None: combined_view_heatmaps[view_name] = _normalize_map(heatmap) return { "query_outputs": query_outputs, "reference_outputs": reference_outputs, **breakdown, "branch_heatmaps": branch_heatmaps, "combined_view_heatmaps": combined_view_heatmaps, } finally: model.train(was_training)