Method forward

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

(self, x: Tensor)

        # RoPE does not need init here, it's handled internally

    def forward(self, x: Tensor) -> Tensor:
        B, _, H, W = x.shape

        # 1. Project in: (B, C, H, W) -> (B, D, H, W)
        x = self.proj_in(x)

        # 2. Reshape for transformer: (B, D, H, W) -> (B, H*W, D)
        x = x.flatten(2).transpose(1, 2)

        # 3. Get RoPE:
        rope_sincos = self.rope_embed(H=H, W=W)

        # 4. Transformer blocks
        for blk in self.blocks:
            x = blk(x, rope_sincos)

        # 5. Final Norm
        x = self.norm(x)

        # 6. Reshape back to image-like: (B, H*W, D) -> (B, D, H, W)
        x = x.transpose(1, 2).reshape(B, self.embed_dim, H, W)

        # 7. Project out: (B, D, H, W) -> (B, C_out * up_factor^2, H, W)
        x = self.proj_out(x)

        # 8. Pixel Shuffle: (B, C_out * up_factor^2, H, W) -> (B, C_out, H*up_factor, W*up_factor)
        x = self.pixel_shuffle(x)

        return x


# Factory functions for different model sizes