Class DinoV3PixelDecoder

vtp/models/decoders/pixel_decoder.py:15–162

DINOv3-based Pixel Decoder for image reconstruction. This decoder uses a transformer architecture with RoPE position embeddings and PixelShuffle for upsampling latent features back to pixel space.

class DinoV3PixelDecoder(nn.Module):
    """DINOv3-based Pixel Decoder for image reconstruction.

    This decoder uses a transformer architecture with RoPE position embeddings
    and PixelShuffle for upsampling latent features back to pixel space.
    """

    def __init__(
        self,
        *,
        in_chans: int = 256,
        out_chans: int = 3,
        upscale_factor: int = 16,
        # ViT params
        pos_embed_rope_base: float = 100.0,
        pos_embed_rope_min_period: Optional[float] = None,
        pos_embed_rope_max_period: Optional[float] = None,
        pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate",
        pos_embed_rope_shift_coords: Optional[float] = None,
        pos_embed_rope_jitter_coords: Optional[float] = None,
        pos_embed_rope_rescale_coords: Optional[float] = None,
        pos_embed_rope_dtype: str = "bf16",
        embed_dim: int = 1024,
        depth: int = 24,
        num_heads: int = 16,
        ffn_ratio: float = 4.0,
        qkv_bias: bool = True,
        drop_path_rate: float = 0.0,
        layerscale_init: Optional[float] = None,
        norm_layer: str = "layernorm",
        ffn_layer: str = "swiglu",
        ffn_bias: bool = True,
        proj_bias: bool = True,
        mask_k_bias: bool = False,
        device: Optional[Any] = None,
        use_qk_norm: bool = False,
        **ignored_kwargs,
    ):
        super().__init__()
        if len(ignored_kwargs) > 0:
            logger.warning(f"Ignored kwargs: {ignored_kwargs}")
        del ignored_kwargs

        norm_layer_cls = norm_layer_dict[norm_layer]
        self.embed_dim = embed_dim
        self.num_heads = num_heads

        # 1. Input projection
        self.proj_in = nn.Conv2d(
            in_chans, embed_dim, kernel_size=1, bias=proj_bias
        )

        # 2. RoPE
        self.rope_embed = RopePositionEmbedding(
            embed_dim=embed_dim,
            num_heads=num_heads,
            base=pos_embed_rope_base,
            min_period=pos_embed_rope_min_period,