Release SKING_DDJ_v104 production inference bundle
f02dec9 verified about 2 months ago
37.6 kB
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from SkingToolkit.dense_uv_parser.uv_topology import (
    build_outer_uv_graph,
    build_simple_uv_topology,
)


def norm_groups(channels):
    for groups in (8, 4, 2, 1):
        if channels % groups == 0:
            return groups
    return 1


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.GroupNorm(norm_groups(out_channels), out_channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.GroupNorm(norm_groups(out_channels), out_channels),
            nn.SiLU(inplace=True),
        )

    def forward(self, x):
        return self.block(x)


class DownBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.down = nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
        self.block = ConvBlock(out_channels, out_channels)

    def forward(self, x):
        return self.block(self.down(x))


class UpBlock(nn.Module):
    def __init__(self, in_channels, skip_channels, out_channels):
        super().__init__()
        self.block = ConvBlock(in_channels + skip_channels, out_channels)

    def forward(self, x, skip):
        x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False)
        return self.block(torch.cat([x, skip], dim=1))


class MultiViewSemanticFusion(nn.Module):
    """Fuse frozen per-view semantic features into parser bottleneck features."""

    def __init__(
        self,
        raw_feature_dim,
        semantic_channels,
        bottleneck_channels,
        view_classes,
        attention_heads=4,
        layers=1,
        dropout=0.05,
    ):
        super().__init__()
        if raw_feature_dim < 1 or semantic_channels < 1:
            raise ValueError("Semantic feature dimensions must be positive.")
        if view_classes < 1:
            raise ValueError("Semantic fusion requires view-conditioned parser inputs.")
        if attention_heads < 1:
            raise ValueError("semantic_attention_heads must be positive.")
        if semantic_channels % attention_heads != 0:
            raise ValueError("semantic_channels must be divisible by attention_heads.")
        if layers < 1:
            raise ValueError("semantic_layers must be positive.")
        if not 0.0 <= dropout < 1.0:
            raise ValueError("semantic_dropout must be in [0, 1).")
        self.raw_feature_dim = int(raw_feature_dim)
        self.semantic_channels = int(semantic_channels)
        self.view_classes = int(view_classes)
        self.input_projection = nn.Sequential(
            nn.LayerNorm(raw_feature_dim),
            nn.Linear(raw_feature_dim, semantic_channels),
            nn.GELU(),
        )
        self.view_embedding = nn.Parameter(
            torch.randn(view_classes, semantic_channels) * 0.02
        )
        self.encoder = nn.ModuleList(
            nn.TransformerEncoderLayer(
                d_model=semantic_channels,
                nhead=attention_heads,
                dim_feedforward=semantic_channels * 4,
                dropout=dropout,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            for _ in range(layers)
        )
        self.modulation = nn.Sequential(
            nn.LayerNorm(semantic_channels * 2),
            nn.Linear(semantic_channels * 2, semantic_channels * 2),
            nn.GELU(),
            nn.Linear(semantic_channels * 2, bottleneck_channels * 2),
        )
        self.summary = nn.Sequential(
            nn.LayerNorm(semantic_channels),
            nn.Linear(semantic_channels, semantic_channels),
            nn.GELU(),
        )
        # Preserve the geometry-only initialization. The adapter starts by
        # contributing no FiLM shift/scale and learns semantic corrections.
        nn.init.zeros_(self.modulation[-1].weight)
        nn.init.zeros_(self.modulation[-1].bias)

    def forward(self, raw_features, view_ids, sample_count):
        if raw_features.dim() == 3:
            # Training may cache multiple complete view groups per skin, for
            # example the inference pair followed by a privileged right-side
            # pair. Each group is still fused independently below.
            raw_features = raw_features.reshape(-1, raw_features.shape[-1])
        if raw_features.dim() != 2 or raw_features.shape != (
            sample_count,
            self.raw_feature_dim,
        ):
            raise ValueError(
                "Semantic features must be shaped NxD or BxVxD; got "
                f"{tuple(raw_features.shape)}."
            )
        if sample_count % self.view_classes != 0:
            raise ValueError(
                f"Semantic sample count {sample_count} is not divisible by "
                f"view count {self.view_classes}."
            )
        grouped_view_ids = view_ids.reshape(-1, self.view_classes)
        expected_ids = torch.arange(
            self.view_classes, device=view_ids.device
        ).view(1, -1)
        if not torch.equal(grouped_view_ids, expected_ids.expand_as(grouped_view_ids)):
            raise ValueError("Semantic fusion requires canonical grouped view order.")

        batch = sample_count // self.view_classes
        tokens = self.input_projection(raw_features.float()).reshape(
            batch, self.view_classes, self.semantic_channels
        )
        tokens = tokens + self.view_embedding.unsqueeze(0)
        for layer in self.encoder:
            tokens = layer(tokens)
        pooled = tokens.mean(dim=1)
        per_view = torch.cat(
            [tokens, pooled.unsqueeze(1).expand_as(tokens)], dim=-1
        ).reshape(sample_count, self.semantic_channels * 2)
        return self.modulation(per_view), self.summary(pooled)


class SpatialSemanticFusion(nn.Module):
    """Project frozen 2D semantic features into the parser bottleneck.
    The final projection starts at zero, preserving the stable geometry parser
    at initialization while allowing spatial semantic corrections to emerge
    during training.
    """

    def __init__(
        self,
        raw_feature_dim,
        semantic_channels,
        bottleneck_channels,
    ):
        super().__init__()
        if raw_feature_dim < 1 or semantic_channels < 1:
            raise ValueError("Spatial semantic feature dimensions must be positive.")
        self.raw_feature_dim = int(raw_feature_dim)
        self.semantic_channels = int(semantic_channels)
        self.input_norm = nn.LayerNorm(raw_feature_dim)
        self.input_projection = nn.Conv2d(
            raw_feature_dim, semantic_channels, kernel_size=1
        )
        self.activation = nn.GELU()
        self.output_projection = nn.Conv2d(
            semantic_channels, bottleneck_channels, kernel_size=1
        )
        nn.init.zeros_(self.output_projection.weight)
        nn.init.zeros_(self.output_projection.bias)

    def forward(self, raw_features, sample_count, output_size):
        if raw_features.dim() == 5:
            raw_features = raw_features.reshape(
                -1,
                raw_features.shape[-3],
                raw_features.shape[-2],
                raw_features.shape[-1],
            )
        if (
            raw_features.dim() != 4
            or raw_features.shape[0] != sample_count
            or raw_features.shape[1] != self.raw_feature_dim
        ):
            raise ValueError(
                "Spatial semantic features must be shaped NxCxHxW or "
                f"BxVxCxHxW; got {tuple(raw_features.shape)}."
            )
        normalized = self.input_norm(
            raw_features.float().permute(0, 2, 3, 1)
        ).permute(0, 3, 1, 2)
        residual = self.output_projection(
            self.activation(self.input_projection(normalized))
        )
        return F.interpolate(
            residual,
            size=output_size,
            mode="bilinear",
            align_corners=False,
        )


class OuterUVGraphBlock(nn.Module):
    """Message passing over physical neighbours of Minecraft outer texels."""

    def __init__(self, channels, dropout=0.05):
        super().__init__()
        self.self_projection = nn.Linear(channels, channels)
        self.neighbour_projection = nn.Linear(channels, channels)
        self.norm = nn.LayerNorm(channels)
        self.dropout = nn.Dropout(dropout)

    def forward(self, nodes, edge_index, degree):
        source, target = edge_index
        neighbour_sum = torch.zeros_like(nodes)
        neighbour_sum.index_add_(1, target, nodes[:, source])
        neighbour_mean = neighbour_sum / degree.view(1, -1, 1).clamp_min(1.0)
        update = self.self_projection(nodes)
        update = update + self.neighbour_projection(neighbour_mean)
        return self.norm(nodes + self.dropout(F.gelu(update)))


class ProjectedOuterUVTopologyHead(nn.Module):
    """Predict outer alpha from projected image features and cube topology."""

    def __init__(
        self,
        input_channels,
        global_context_dim,
        hidden_channels=64,
        layers=3,
        dropout=0.05,
        uv_size=64,
    ):
        super().__init__()
        if layers < 1:
            raise ValueError("outer_uv_topology_layers must be positive.")
        flat_indices, edge_index = build_outer_uv_graph()
        topology = build_simple_uv_topology()
        parts = topology.part.reshape(-1)[flat_indices]
        faces = topology.face.reshape(-1)[flat_indices]
        local_uv = topology.local_uv.reshape(-1, 2)[flat_indices]
        degree = torch.bincount(
            edge_index[1],
            minlength=flat_indices.numel(),
        ).float()
        self.uv_size = int(uv_size)
        self.register_buffer("flat_indices", flat_indices, persistent=False)
        self.register_buffer("edge_index", edge_index, persistent=False)
        self.register_buffer("degree", degree, persistent=False)
        self.register_buffer("parts", parts, persistent=False)
        self.register_buffer("faces", faces, persistent=False)
        self.register_buffer("local_uv", local_uv, persistent=False)
        self.input_projection = nn.Linear(input_channels, hidden_channels)
        self.part_embedding = nn.Embedding(6, hidden_channels)
        self.face_embedding = nn.Embedding(6, hidden_channels)
        self.position_projection = nn.Linear(2, hidden_channels)
        self.global_projection = nn.Sequential(
            nn.LayerNorm(global_context_dim),
            nn.Linear(global_context_dim, hidden_channels),
            nn.GELU(),
        )
        self.blocks = nn.ModuleList(
            OuterUVGraphBlock(hidden_channels, dropout=dropout)
            for _ in range(layers)
        )
        self.output = nn.Sequential(
            nn.LayerNorm(hidden_channels),
            nn.Linear(hidden_channels, 1),
        )
        # A very negative initialization made the one-epoch recipe prone to
        # the all-transparent solution after precision penalties were added.
        # Keep a sparse prior without starving early positive gradients.
        nn.init.constant_(self.output[-1].bias, -1.0)

    def forward(self, atlas_features, global_context):
        if atlas_features.dim() != 4:
            raise ValueError("Projected outer UV features must be BCHW.")
        if atlas_features.shape[-2:] != (self.uv_size, self.uv_size):
            raise ValueError(
                "Projected outer UV feature size must match the UV atlas."
            )
        flattened = atlas_features.flatten(2).transpose(1, 2)
        nodes = flattened.index_select(1, self.flat_indices)
        nodes = self.input_projection(nodes.float())
        nodes = nodes + self.part_embedding(self.parts).unsqueeze(0)
        nodes = nodes + self.face_embedding(self.faces).unsqueeze(0)
        nodes = nodes + self.position_projection(self.local_uv).unsqueeze(0)
        nodes = nodes + self.global_projection(global_context.float()).unsqueeze(1)
        for block in self.blocks:
            nodes = block(nodes, self.edge_index, self.degree)
        node_logits = self.output(nodes).squeeze(-1)
        atlas_logits = node_logits.new_full(
            (node_logits.shape[0], self.uv_size * self.uv_size),
            -12.0,
        )
        atlas_logits[:, self.flat_indices] = node_logits
        return atlas_logits.reshape(
            node_logits.shape[0],
            1,
            self.uv_size,
            self.uv_size,
        )


class DenseUVParserNet(nn.Module):
    """Predict dense Minecraft UV routing for each render pixel."""

    def __init__(
        self,
        input_channels=4,
        base_channels=32,
        part_classes=6,
        face_classes=6,
        layer_classes=None,
        layer_face_classes=12,
        uv_size=64,
        uv_classification=True,
        view_classes=0,
        predict_affine=False,
        affine_translation_scale=0.0,
        affine_scale_range=0.0,
        surface_classes=0,
        geometry_only=False,
        feature_dropout=0.0,
        semantic_feature_dim=0,
        semantic_channels=128,
        semantic_attention_heads=4,
        semantic_layers=1,
        semantic_dropout=0.05,
        semantic_spatial_feature_dim=0,
        semantic_spatial_channels=64,
        predict_confidence=False,
        route_role_spatial_prior=False,
        route_prior_height=32,
        route_prior_width=16,
        route_prior_logit_cap=1.5,
        route_prior_dropout=0.10,
        predict_head_accessories=False,
        predict_hat_components=False,
        predict_head_ownership=False,
        predict_headwear=False,
        predict_head_semantics=False,
        predict_head_surface=False,
        predict_headwear_presence=False,
        predict_headphone_presence=False,
        accessory_route_threshold=0.90,
        predict_outer_uv_occupancy=False,
        outer_uv_feature_channels=32,
        outer_uv_topology_channels=64,
        outer_uv_topology_layers=3,
        outer_uv_topology_dropout=0.05,
        outer_uv_route_evidence_dropout=1.0,
    ):
        super().__init__()
        self.accessory_route_threshold = float(accessory_route_threshold)
        if not 0.5 <= self.accessory_route_threshold <= 1.:
            raise ValueError("Accessory seed confidence must lie in [0.5, 1]")
        self.predict_hat_components = bool(predict_hat_components)
        self.predict_head_accessories = bool(predict_head_accessories)
        if self.predict_head_accessories:
            from SkingToolkit.dense_uv_parser.accessories import HeadAccessoryHead
            self.accessory_head = HeadAccessoryHead(semantic_spatial_feature_dim, self.predict_hat_components)
        self.predict_headphone_presence = bool(predict_headphone_presence)
        self.predict_headwear_presence = bool(predict_headwear_presence)
        self.predict_headwear = bool(predict_headwear)
        if self.predict_headwear:
            from SkingToolkit.dense_uv_parser.headwear import HeadwearHead
            self.headwear_head = HeadwearHead(semantic_spatial_feature_dim,self.predict_headwear_presence)
        self.predict_head_ownership = bool(predict_head_ownership)
        if self.predict_head_ownership:
            from SkingToolkit.dense_uv_parser.ownership import HeadOwnershipHead
            self.ownership_head = HeadOwnershipHead(semantic_spatial_feature_dim,self.predict_headphone_presence)
        self.predict_head_semantics = bool(predict_head_semantics)
        self.predict_head_surface = bool(predict_head_surface)
        if self.predict_head_semantics:
            from SkingToolkit.dense_uv_parser.head_semantics import HeadSemanticsHead
            self.head_semantics_head = HeadSemanticsHead(semantic_spatial_feature_dim,predict_surface=self.predict_head_surface)
        self.geometry_only = bool(geometry_only)
        if layer_classes is None:
            layer_classes = 3 if self.geometry_only else 2
        self.layer_classes = int(layer_classes)
        self.uv_size = int(uv_size)
        self.uv_classification = bool(uv_classification) and not self.geometry_only
        self.view_classes = int(view_classes)
        self.predict_affine = bool(predict_affine)
        self.affine_translation_scale = float(affine_translation_scale)
        self.affine_scale_range = float(affine_scale_range)
        self.surface_classes = int(surface_classes)
        self.semantic_feature_dim = int(semantic_feature_dim)
        self.semantic_channels = int(semantic_channels)
        self.semantic_attention_heads = int(semantic_attention_heads)
        self.semantic_layers = int(semantic_layers)
        self.semantic_dropout = float(semantic_dropout)
        self.semantic_spatial_feature_dim = int(semantic_spatial_feature_dim)
        self.semantic_spatial_channels = int(semantic_spatial_channels)
        self.predict_confidence = bool(predict_confidence)
        self.route_role_spatial_prior = bool(route_role_spatial_prior)
        self.route_prior_height = int(route_prior_height)
        self.route_prior_width = int(route_prior_width)
        self.route_prior_logit_cap = float(route_prior_logit_cap)
        self.route_prior_dropout = float(route_prior_dropout)
        self.predict_outer_uv_occupancy = bool(predict_outer_uv_occupancy)
        self.outer_uv_feature_channels = int(outer_uv_feature_channels)
        self.outer_uv_topology_channels = int(outer_uv_topology_channels)
        self.outer_uv_topology_layers = int(outer_uv_topology_layers)
        self.outer_uv_topology_dropout = float(outer_uv_topology_dropout)
        self.outer_uv_route_evidence_dropout = float(
            outer_uv_route_evidence_dropout
        )
        if self.route_prior_height < 1 or self.route_prior_width < 1:
            raise ValueError("Route-prior dimensions must be positive.")
        if self.route_prior_logit_cap <= 0.0:
            raise ValueError("route_prior_logit_cap must be positive.")
        if not 0.0 <= self.route_prior_dropout < 1.0:
            raise ValueError("route_prior_dropout must be in [0, 1).")
        if self.route_role_spatial_prior and (
            not self.geometry_only or self.view_classes < 1
        ):
            raise ValueError(
                "The fixed-view route-role prior requires geometry_only with view classes."
            )
        self.feature_dropout_probability = float(feature_dropout)
        if not 0.0 <= self.feature_dropout_probability < 1.0:
            raise ValueError("feature_dropout must be in [0, 1).")
        if not 0.0 <= self.outer_uv_route_evidence_dropout <= 1.0:
            raise ValueError(
                "outer_uv_route_evidence_dropout must be in [0, 1]."
            )
        c = base_channels
        self.stem = ConvBlock(input_channels + self.view_classes, c)
        self.down1 = DownBlock(c, c * 2)
        self.down2 = DownBlock(c * 2, c * 4)
        self.down3 = DownBlock(c * 4, c * 8)
        self.mid = ConvBlock(c * 8, c * 8)
        self.semantic_fusion = (
            MultiViewSemanticFusion(
                self.semantic_feature_dim,
                self.semantic_channels,
                c * 8,
                self.view_classes,
                attention_heads=self.semantic_attention_heads,
                layers=self.semantic_layers,
                dropout=self.semantic_dropout,
            )
            if self.semantic_feature_dim > 0
            else None
        )
        self.semantic_spatial_fusion = (
            SpatialSemanticFusion(
                self.semantic_spatial_feature_dim,
                self.semantic_spatial_channels,
                c * 8,
            )
            if self.semantic_spatial_feature_dim > 0
            else None
        )
        self.up2 = UpBlock(c * 8, c * 4, c * 4)
        self.up1 = UpBlock(c * 4, c * 2, c * 2)
        self.up0 = UpBlock(c * 2, c, c)
        self.features = nn.Sequential(
            nn.Conv2d(c, c, kernel_size=3, padding=1),
            nn.SiLU(inplace=True),
        )
        self.feature_dropout = nn.Dropout2d(self.feature_dropout_probability)
        self.foreground = nn.Conv2d(c, 1, kernel_size=1)
        self.layer = nn.Conv2d(c, self.layer_classes, kernel_size=1)
        self.route_role_prior = (
            nn.Parameter(
                torch.zeros(
                    self.view_classes,
                    self.layer_classes,
                    self.route_prior_height,
                    self.route_prior_width,
                )
            )
            if self.route_role_spatial_prior
            else None
        )
        self.route_confidence = (
            nn.Conv2d(c, 1, kernel_size=1) if self.predict_confidence else None
        )
        if self.semantic_fusion is not None:
            self.outer_presence_head = nn.Linear(self.semantic_channels, 6)
            self.outer_coverage_head = nn.Linear(self.semantic_channels, 6)
        if self.predict_outer_uv_occupancy:
            if self.view_classes < 1:
                raise ValueError(
                    "Outer UV occupancy prediction requires grouped fixed-view inputs."
                )
            if self.outer_uv_feature_channels < 1:
                raise ValueError("outer_uv_feature_channels must be positive.")
            occupancy_context_channels = c * 8 + (
                self.semantic_channels if self.semantic_fusion is not None else 0
            )
            self.outer_uv_feature_projection = nn.Sequential(
                nn.Conv2d(c, self.outer_uv_feature_channels, kernel_size=1),
                nn.GELU(),
            )
            # Projected feature mean + mean/max p_outer + foreground coverage
            # + number of supporting views.
            occupancy_input_channels = self.outer_uv_feature_channels + 4
            self.outer_uv_occupancy_head = ProjectedOuterUVTopologyHead(
                occupancy_input_channels,
                occupancy_context_channels,
                hidden_channels=self.outer_uv_topology_channels,
                layers=self.outer_uv_topology_layers,
                dropout=self.outer_uv_topology_dropout,
                uv_size=uv_size,
            )
        if not self.geometry_only:
            self.part = nn.Conv2d(c, part_classes, kernel_size=1)
            self.face = nn.Conv2d(c, face_classes, kernel_size=1)
            self.layer_face = (
                nn.Conv2d(c, layer_face_classes, kernel_size=1)
                if layer_face_classes > 0
                else None
            )
            self.uv = nn.Conv2d(c, 2, kernel_size=1)
            if self.uv_classification:
                self.uv_x = nn.Conv2d(c, uv_size, kernel_size=1)
                self.uv_y = nn.Conv2d(c, uv_size, kernel_size=1)
        else:
            self.layer_face = None
        if self.predict_affine:
            if not self.geometry_only and self.surface_classes < 2:
                raise ValueError("Global-affine routing requires at least two static surface classes.")
            if self.surface_classes == 1:
                raise ValueError("Surface routing requires at least two static surface classes.")
            if self.surface_classes > 0:
                self.surface = nn.Conv2d(c, self.surface_classes, kernel_size=1)
            hidden = max(c * 4, 32)
            self.affine_head = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Flatten(),
                nn.Linear(c * 8, hidden),
                nn.SiLU(inplace=True),
                nn.Linear(hidden, 3),
            )
            # Start from the canonical pose; the residual head learns only the
            # small global transform introduced by the configured augmentation.
            nn.init.zeros_(self.affine_head[-1].weight)
            nn.init.zeros_(self.affine_head[-1].bias)

            self.affine_translation_limit = 2.0 * self.affine_translation_scale
            lower_log_scale = math.log(max(1.0 - self.affine_scale_range, 1e-6))
            upper_log_scale = math.log1p(self.affine_scale_range)
            self.affine_log_scale_limit = max(abs(lower_log_scale), abs(upper_log_scale))

    def _runtime_semantic_features(self, images, foreground=None):
        backbone = getattr(self, "_runtime_semantic_backbone", None)
        if backbone is None:
            raise ValueError(
                "This parser requires semantic_features or an attached semantic "
                "runtime backbone."
            )
        rgb = images[:, :3]
        if foreground is not None:
            if foreground.dim() == 3:
                foreground = foreground.unsqueeze(1)
            if (
                foreground.dim() != 4
                or foreground.shape[0] != rgb.shape[0]
                or foreground.shape[1] != 1
            ):
                raise ValueError(
                    "semantic_foreground must be shaped Nx1xHxW or NxHxW; got "
                    f"{tuple(foreground.shape)}."
                )
            if foreground.shape[-2:] != rgb.shape[-2:]:
                foreground = F.interpolate(
                    foreground.float(),
                    size=rgb.shape[-2:],
                    mode="nearest",
                )
            foreground = foreground.to(device=rgb.device, dtype=rgb.dtype).clamp(0, 1)
            # A fixed neutral background prevents the frozen spatial tower from
            # learning accidental correlations with randomized training colors.
            rgb = rgb * foreground + 0.5 * (1.0 - foreground)
        with torch.no_grad():
            if hasattr(backbone, "encode_dense"):
                return backbone.encode_dense(rgb)
            return backbone.encode_global(rgb)

    def forward(
        self,
        x,
        view_ids=None,
        semantic_features=None,
        semantic_foreground=None,
    ):
        source_images = x
        if self.view_classes > 0:
            if view_ids is None:
                raise ValueError("view_ids are required for a view-conditioned dense UV parser.")
            if view_ids.shape != (x.shape[0],):
                raise ValueError(f"Expected view_ids shape {(x.shape[0],)}, got {tuple(view_ids.shape)}.")
            if view_ids.min() < 0 or view_ids.max() >= self.view_classes:
                raise ValueError(f"view_ids must be in [0, {self.view_classes - 1}].")
            view_one_hot = F.one_hot(view_ids.long(), num_classes=self.view_classes).to(dtype=x.dtype)
            view_one_hot = view_one_hot.view(x.shape[0], self.view_classes, 1, 1)
            x = torch.cat([x, view_one_hot.expand(-1, -1, x.shape[2], x.shape[3])], dim=1)

        s0 = self.stem(x)
        s1 = self.down1(s0)
        s2 = self.down2(s1)
        s3 = self.down3(s2)
        x = self.mid(s3)
        if (
            self.semantic_fusion is not None
            or self.semantic_spatial_fusion is not None
        ) and semantic_features is None:
            semantic_features = self._runtime_semantic_features(
                source_images, foreground=semantic_foreground
            )
        if isinstance(semantic_features, dict):
            semantic_global = semantic_features.get(
                "raw_global", semantic_features.get("global")
            )
            semantic_spatial = semantic_features.get(
                "raw_spatial", semantic_features.get("spatial")
            )
        else:
            semantic_global = semantic_features
            semantic_spatial = None
        if self.semantic_spatial_fusion is not None:
            if semantic_spatial is None:
                raise ValueError(
                    "This parser checkpoint requires spatial semantic features."
                )
            x = x + self.semantic_spatial_fusion(
                semantic_spatial,
                source_images.shape[0],
                x.shape[-2:],
            ).to(dtype=x.dtype)
        visual_summary = x.mean(dim=(2, 3))
        if self.predict_outer_uv_occupancy:
            if visual_summary.shape[0] % self.view_classes != 0:
                raise ValueError(
                    "Outer UV occupancy prediction requires complete fixed-view groups."
                )
            visual_summary = visual_summary.reshape(
                -1, self.view_classes, visual_summary.shape[-1]
            ).mean(dim=1)
        semantic_summary = None
        if self.semantic_fusion is not None:
            if semantic_global is None:
                raise ValueError(
                    "This parser checkpoint requires global semantic features."
                )
            modulation, semantic_summary = self.semantic_fusion(
                semantic_global,
                view_ids,
                source_images.shape[0],
            )
            scale, shift = modulation.chunk(2, dim=1)
            x = x * (1.0 + scale.unsqueeze(-1).unsqueeze(-1))
            x = x + shift.unsqueeze(-1).unsqueeze(-1)
        affine = None
        if self.predict_affine:
            raw_affine = torch.tanh(self.affine_head(x))
            affine = torch.stack(
                [
                    raw_affine[:, 0] * self.affine_translation_limit,
                    raw_affine[:, 1] * self.affine_translation_limit,
                    raw_affine[:, 2] * self.affine_log_scale_limit,
                ],
                dim=1,
            )
        x = self.up2(x, s2)
        x = self.up1(x, s1)
        x = self.up0(x, s0)
        x = self.features(x)
        occupancy_feature_source = x
        x = self.feature_dropout(x)
        layer_evidence = self.layer(x)
        outputs = {
            "foreground": self.foreground(x),
            "layer": layer_evidence,
        }
        if self.route_role_prior is not None:
            selected_prior_raw = self.route_role_prior.index_select(
                0, view_ids.long()
            )
            selected_prior = self.route_prior_logit_cap * torch.tanh(
                selected_prior_raw / self.route_prior_logit_cap
            )
            if self.training and self.route_prior_dropout > 0.0:
                keep = (
                    torch.rand(
                        selected_prior.shape[0],
                        1,
                        1,
                        1,
                        device=selected_prior.device,
                    )
                    >= self.route_prior_dropout
                )
                selected_prior = selected_prior * keep.to(
                    dtype=selected_prior.dtype
                )
            route_prior = F.interpolate(
                selected_prior,
                size=layer_evidence.shape[-2:],
                mode="bilinear",
                align_corners=False,
            ).to(dtype=layer_evidence.dtype)
            outputs["layer"] = layer_evidence + route_prior
            outputs["route_role_evidence"] = layer_evidence
            outputs["route_role_prior"] = route_prior
            outputs["route_role_prior_raw"] = self.route_role_prior
        if self.route_confidence is not None:
            outputs["route_confidence"] = self.route_confidence(x)
        if semantic_summary is not None:
            outputs["outer_presence_logits"] = self.outer_presence_head(
                semantic_summary
            )
            outputs["outer_coverage"] = torch.sigmoid(
                self.outer_coverage_head(semantic_summary)
            )
        if self.predict_outer_uv_occupancy:
            occupancy_summary = (
                torch.cat([visual_summary, semantic_summary], dim=1)
                if semantic_summary is not None
                else visual_summary
            )
            # Occupancy learns from the parser representation without allowing
            # its sparse atlas loss to destabilize the proven image-space route
            # trunk. A separate agreement loss later transfers calibrated UV
            # structure back into the route head after warm-up.
            outputs["outer_uv_features"] = self.outer_uv_feature_projection(
                occupancy_feature_source.detach()
            )
            outputs["outer_uv_global_context"] = occupancy_summary.detach()
        if not self.geometry_only:
            outputs["part"] = self.part(x)
            outputs["face"] = self.face(x)
            outputs["uv"] = torch.sigmoid(self.uv(x))
            if self.layer_face is not None:
                outputs["layer_face"] = self.layer_face(x)
            if self.uv_classification:
                outputs["uv_x"] = self.uv_x(x)
                outputs["uv_y"] = self.uv_y(x)
        if affine is not None:
            # [tx, ty, log_scale]. tx/ty are affine_grid normalized coordinates.
            outputs["affine"] = affine
            if self.surface_classes > 0:
                outputs["surface"] = self.surface(x)
        if self.predict_head_accessories:
            prediction=self.predict_accessories(source_images,semantic_foreground,return_components=self.predict_hat_components)
            if self.predict_hat_components:
                outputs["accessory_logits"],outputs["hat_component_logits"]=prediction
            else:outputs["accessory_logits"]=prediction
            outputs["accessory_route_threshold"] = self.accessory_route_threshold
        if self.predict_headwear:
            prediction = self.predict_headwear_components(source_images, semantic_foreground,return_presence=self.predict_headwear_presence)
            if self.predict_headwear_presence:outputs["headwear_logits"],outputs["headwear_presence_logits"]=prediction
            else:outputs["headwear_logits"]=prediction
        if self.predict_head_ownership:
            prediction = self.predict_ownership(source_images, semantic_foreground, return_presence=self.predict_headphone_presence)
            if self.predict_headphone_presence:
                outputs["head_ownership_logits"],outputs["headphone_presence_logit"] = prediction
            else:outputs["head_ownership_logits"] = prediction
        if self.predict_head_semantics:
            from SkingToolkit.dense_uv_parser.head_semantics import project_semantics
            joint_result=self.predict_joint_head_semantics(source_images,semantic_foreground,return_surface=self.predict_head_surface)
            joint,presence=joint_result[:2]
            if self.predict_head_surface:outputs['head_surface_logits']=joint_result[2]
            outputs['head_color_ownership_logits'] = outputs['head_ownership_logits']
            outputs['head_semantics_logits'] = joint
            outputs['headwear_logits'] = project_semantics(joint, 'headwear')
            outputs['head_ownership_logits'] = project_semantics(joint, 'ownership')
            outputs['headwear_presence_logits'] = presence
        return outputs

    def predict_joint_head_semantics(self, images, foreground=None, return_surface=False):
        from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits
        crop = head_crop(images, foreground)
        semantic = self._runtime_semantic_features(crop)['raw_spatial']
        result=self.head_semantics_head(crop,semantic,return_surface=return_surface)
        logits,presence=result[:2]
        restored=(restore_logits(logits, images.shape[-2:]),presence)
        return (*restored,restore_logits(result[2],images.shape[-2:])) if return_surface else restored

    def predict_headwear_components(self, images, foreground=None, return_presence=False):
        from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits
        crop = head_crop(images, foreground)
        semantic = self._runtime_semantic_features(crop)["raw_spatial"]
        prediction=self.headwear_head(crop,semantic,return_presence=return_presence)
        if return_presence:return restore_logits(prediction[0],images.shape[-2:]),prediction[1]
        return restore_logits(prediction,images.shape[-2:])

    def predict_ownership(self, images, foreground=None, return_presence=False):
        from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits
        crop = head_crop(images, foreground)
        semantic = self._runtime_semantic_features(crop)["raw_spatial"]
        prediction = self.ownership_head(crop, semantic, return_presence=return_presence)
        if return_presence:return restore_logits(prediction[0], images.shape[-2:]),prediction[1]
        return restore_logits(prediction, images.shape[-2:])

    def predict_accessories(self, images, foreground=None, return_components=False):
        from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits
        crop = head_crop(images, foreground)
        semantic = self._runtime_semantic_features(crop)["raw_spatial"]
        logits = self.accessory_head(crop, semantic, return_components=return_components)
        if isinstance(logits,tuple):return tuple(restore_logits(x,images.shape[-2:]) for x in logits)
        return restore_logits(logits, images.shape[-2:])

    def predict_projected_outer_uv_occupancy(
        self,
        atlas_features,
        global_context,
    ):
        if not self.predict_outer_uv_occupancy:
            raise ValueError("This parser has no outer UV occupancy head.")
        return self.outer_uv_occupancy_head(
            atlas_features,
            global_context,
        )


def count_parameters(model):
    return sum(param.numel() for param in model.parameters() if param.requires_grad)