artist_style_dinov3/explain.py

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)