analogy_projector.py

"""
ADOBE CONFIDENTIAL
Copyright 2024 Adobe
All Rights Reserved.
NOTICE: All information contained herein is, and remains
the property of Adobe and its suppliers, if any. The intellectual
and technical concepts contained herein are proprietary to Adobe
and its suppliers and are protected by all applicable intellectual
property laws, including trade secret and copyright laws.
Dissemination of this information or reproduction of this material
is strictly forbidden unless prior written permission is obtained
from Adobe.
"""

import einops
import numpy as np
import torch as th
import torch.nn as nn
from diffusers import ModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
# REf: https://github.com/tatp22/multidim-positional-encoding/tree/master


OUT_SIZE = 768
IN_SIZE = 2048


def get_emb(sin_inp):
    """
    Gets a base embedding for one dimension with sin and cos intertwined
    """
    emb = th.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
    return th.flatten(emb, -2, -1)


class PositionalEncoding1D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        super(PositionalEncoding1D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 2) * 2)
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("cached_penc", None, persistent=False)

    def forward(self, tensor):
        """
        :param tensor: A 3d tensor of size (batch_size, x, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, ch)
        """
        if len(tensor.shape) != 3:
            raise RuntimeError("The input tensor has to be 3d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, orig_ch = tensor.shape
        pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype)
        sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq)
        emb_x = get_emb(sin_inp_x)
        emb = th.zeros((x, self.channels), device=tensor.device, dtype=tensor.dtype)
        emb[:, : self.channels] = emb_x

        self.cached_penc = emb[None, :, :orig_ch].repeat(batch_size, 1, 1)
        return self.cached_penc



class PositionalEncoding3D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        super(PositionalEncoding3D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 6) * 2)
        if channels % 2:
            channels += 1
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("cached_penc", None, persistent=False)

    def forward(self, tensor):
        """
        :param tensor: A 5d tensor of size (batch_size, x, y, z, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch)
        """
        if len(tensor.shape) != 5:
            raise RuntimeError("The input tensor has to be 5d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, y, z, orig_ch = tensor.shape
        pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype)
        pos_y = th.arange(y, device=tensor.device, dtype=self.inv_freq.dtype)
        pos_z = th.arange(z, device=tensor.device, dtype=self.inv_freq.dtype)
        sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq)
        sin_inp_y = th.einsum("i,j->ij", pos_y, self.inv_freq)
        sin_inp_z = th.einsum("i,j->ij", pos_z, self.inv_freq)
        emb_x = get_emb(sin_inp_x).unsqueeze(1).unsqueeze(1)
        emb_y = get_emb(sin_inp_y).unsqueeze(1)
        emb_z = get_emb(sin_inp_z)
        emb = th.zeros(
            (x, y, z, self.channels * 3),
            device=tensor.device,
            dtype=tensor.dtype,
        )
        emb[:, :, :, : self.channels] = emb_x
        emb[:, :, :, self.channels : 2 * self.channels] = emb_y
        emb[:, :, :, 2 * self.channels :] = emb_z

        self.cached_penc = emb[None, :, :, :, :orig_ch].repeat(batch_size, 1, 1, 1, 1)
        return self.cached_penc

class AnalogyProjector(ModelMixin, ConfigMixin):
    
    @register_to_config
    def __init__(self):
        super(AnalogyProjector, self).__init__()
        self.projector = DinoSiglipMixer()
        self.pos_embd_1D = PositionalEncoding1D(OUT_SIZE)
        self.pos_embd_3D = PositionalEncoding3D(OUT_SIZE)
        

    def forward(self, dino_in, siglip_in, batch_size):
        
        image_embeddings = self.projector(dino_in, siglip_in)
        
        image_embeddings = einops.rearrange(image_embeddings, '(k b) t d -> b k t d', b=batch_size)
        image_embeddings = self.position_embd(image_embeddings)
        return image_embeddings

    def position_embd(self, image_embeddings, concat=False):
        canvas_embd = image_embeddings[:, :, 1:, :]
        batch_size = canvas_embd.shape[0]
        type_size = canvas_embd.shape[1]
        xy_size = canvas_embd.shape[2]
        
        x_size = int(xy_size ** 0.5)

        canvas_embd = canvas_embd.reshape(batch_size, type_size, x_size, x_size, -1)
        if concat:
            canvas_embd = th.cat([canvas_embd, self.pos_embd_3D(canvas_embd)], -1)
        else:
            canvas_embd = self.pos_embd_3D(canvas_embd) + canvas_embd
        canvas_embd = canvas_embd.reshape(batch_size, type_size, xy_size, -1)

        class_embd = image_embeddings[:, :, 0, :]
        if concat:
            class_embd = th.cat([class_embd, self.pos_embd_1D(class_embd)], -1)
        else:    
            class_embd = self.pos_embd_1D(class_embd) + class_embd
        all_embd_list = []
        for i in range(type_size):
            all_embd_list.append(class_embd[:, i:i+1])
            all_embd_list.append(canvas_embd[:, i])
        image_embeddings = th.cat(all_embd_list, 1)
        return image_embeddings


class HighLowMixer(th.nn.Module):
    def __init__(self, in_size=IN_SIZE, out_size=OUT_SIZE):
        super().__init__()
        mid_size = (in_size + out_size) // 2
        
        self.lower_projector = th.nn.Sequential(
            th.nn.LayerNorm(IN_SIZE//2),
            th.nn.SiLU()
        )
        self.upper_projector = th.nn.Sequential(
            th.nn.LayerNorm(IN_SIZE//2),
            th.nn.SiLU()
        )
        self.projectors = th.nn.ModuleList([
            # add layer norm
            th.nn.Linear(in_size, mid_size),
            th.nn.SiLU(),
            th.nn.Linear(mid_size, out_size)
        ])
        # initialize
        for proj in self.projectors:
            if isinstance(proj, th.nn.Linear):
                th.nn.init.xavier_uniform_(proj.weight)
                th.nn.init.zeros_(proj.bias)

    def forward(self, lower_in, upper_in, ):
        # ALso format lower_in
        lower_in = self.lower_projector(lower_in)
        upper_in = self.upper_projector(upper_in)
        x = th.cat([lower_in, upper_in], -1)
        for proj in self.projectors:
            x = proj(x)
        return x

class DinoSiglipMixer(th.nn.Module):
    def __init__(self, in_size=OUT_SIZE * 2, out_size=OUT_SIZE):
        super().__init__()
        self.dino_projector = HighLowMixer()
        self.siglip_projector = HighLowMixer()
        self.projectors = th.nn.Sequential(
            th.nn.SiLU(),
            th.nn.Linear(in_size, out_size),
        )
        # initialize
        for proj in self.projectors:
            if isinstance(proj, th.nn.Linear):
                th.nn.init.xavier_uniform_(proj.weight)
                th.nn.init.zeros_(proj.bias)

    
    def forward(self, dino_in, siglip_in):
        # ALso format lower_in
        lower, upper = th.chunk(dino_in, 2, -1)
        dino_out = self.dino_projector(lower, upper)
        lower, upper = th.chunk(siglip_in, 2, -1)
        siglip_out = self.siglip_projector(lower, upper)
        x = th.cat([dino_out, siglip_out], -1)
        for proj in self.projectors:
            x = proj(x)
        return x