clip/vitseg.py

import math
from posixpath import basename, dirname, join

# import clip
from clip.model import convert_weights
import torch
import json
from torch import nn
from torch.nn import functional as nnf
from torch.nn.modules import activation
from torch.nn.modules.activation import ReLU
from torchvision import transforms

normalize = transforms.Normalize(
    mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)
)

from torchvision.models import ResNet


def process_prompts(conditional, prompt_list, conditional_map):
    # DEPRECATED

    # randomly sample a synonym
    words = [conditional_map[int(i)] for i in conditional]
    words = [
        syns[torch.multinomial(torch.ones(len(syns)), 1, replacement=True).item()]
        for syns in words
    ]
    words = [w.replace("_", " ") for w in words]

    if prompt_list is not None:
        prompt_indices = torch.multinomial(
            torch.ones(len(prompt_list)), len(words), replacement=True
        )
        prompts = [prompt_list[i] for i in prompt_indices]
    else:
        prompts = ["a photo of {}"] * (len(words))

    return [promt.format(w) for promt, w in zip(prompts, words)]


class VITDenseBase(nn.Module):
    def rescaled_pos_emb(self, new_size):
        assert len(new_size) == 2

        a = self.model.positional_embedding[1:].T.view(1, 768, *self.token_shape)
        b = (
            nnf.interpolate(a, new_size, mode="bicubic", align_corners=False)
            .squeeze(0)
            .view(768, new_size[0] * new_size[1])
            .T
        )
        return torch.cat([self.model.positional_embedding[:1], b])

    def visual_forward(self, x_inp, extract_layers=(), skip=False, mask=None):
        with torch.no_grad():
            x_inp = nnf.interpolate(x_inp, (384, 384))

            x = self.model.patch_embed(x_inp)
            cls_token = self.model.cls_token.expand(
                x.shape[0], -1, -1
            )  # stole cls_tokens impl from Phil Wang, thanks
            if self.model.dist_token is None:
                x = torch.cat((cls_token, x), dim=1)
            else:
                x = torch.cat(
                    (cls_token, self.model.dist_token.expand(x.shape[0], -1, -1), x),
                    dim=1,
                )
            x = self.model.pos_drop(x + self.model.pos_embed)

            activations = []
            for i, block in enumerate(self.model.blocks):
                x = block(x)

                if i in extract_layers:
                    # permute to be compatible with CLIP
                    activations += [x.permute(1, 0, 2)]

            x = self.model.norm(x)
            x = self.model.head(self.model.pre_logits(x[:, 0]))

            # again for CLIP compatibility
            # x = x.permute(1, 0, 2)

        return x, activations, None

    def sample_prompts(self, words, prompt_list=None):
        prompt_list = prompt_list if prompt_list is not None else self.prompt_list

        prompt_indices = torch.multinomial(
            torch.ones(len(prompt_list)), len(words), replacement=True
        )
        prompts = [prompt_list[i] for i in prompt_indices]
        return [promt.format(w) for promt, w in zip(prompts, words)]

    def get_cond_vec(self, conditional, batch_size):
        # compute conditional from a single string
        if conditional is not None and type(conditional) == str:
            cond = self.compute_conditional(conditional)
            cond = cond.repeat(batch_size, 1)

        # compute conditional from string list/tuple
        elif (
            conditional is not None
            and type(conditional) in {list, tuple}
            and type(conditional[0]) == str
        ):
            assert len(conditional) == batch_size
            cond = self.compute_conditional(conditional)

        # use conditional directly
        elif (
            conditional is not None
            and type(conditional) == torch.Tensor
            and conditional.ndim == 2
        ):
            cond = conditional

        # compute conditional from image
        elif conditional is not None and type(conditional) == torch.Tensor:
            with torch.no_grad():
                cond, _, _ = self.visual_forward(conditional)
        else:
            raise ValueError("invalid conditional")
        return cond

    def compute_conditional(self, conditional):
        import clip

        dev = next(self.parameters()).device

        if type(conditional) in {list, tuple}:
            text_tokens = clip.tokenize(conditional).to(dev)
            cond = self.clip_model.encode_text(text_tokens)
        else:
            if conditional in self.precomputed_prompts:
                cond = self.precomputed_prompts[conditional].float().to(dev)
            else:
                text_tokens = clip.tokenize([conditional]).to(dev)
                cond = self.clip_model.encode_text(text_tokens)[0]

        return cond


class VITDensePredT(VITDenseBase):
    def __init__(
        self,
        extract_layers=(3, 6, 9),
        cond_layer=0,
        reduce_dim=128,
        n_heads=4,
        prompt="fixed",
        depth=3,
        extra_blocks=0,
        reduce_cond=None,
        fix_shift=False,
        learn_trans_conv_only=False,
        refine=None,
        limit_to_clip_only=False,
        upsample=False,
        add_calibration=False,
        process_cond=None,
        not_pretrained=False,
    ):
        super().__init__()
        # device = 'cpu'

        self.extract_layers = extract_layers
        self.cond_layer = cond_layer
        self.limit_to_clip_only = limit_to_clip_only
        self.process_cond = None

        if add_calibration:
            self.calibration_conds = 1

        self.upsample_proj = (
            nn.Conv2d(reduce_dim, 1, kernel_size=1) if upsample else None
        )

        self.add_activation1 = True

        import timm

        self.model = timm.create_model("vit_base_patch16_384", pretrained=True)
        self.model.head = nn.Linear(768, 512 if reduce_cond is None else reduce_cond)

        for p in self.model.parameters():
            p.requires_grad_(False)

        import clip

        self.clip_model, _ = clip.load("ViT-B/16", device="cpu", jit=False)
        # del self.clip_model.visual

        self.token_shape = (14, 14)

        # conditional
        if reduce_cond is not None:
            self.reduce_cond = nn.Linear(512, reduce_cond)
            for p in self.reduce_cond.parameters():
                p.requires_grad_(False)
        else:
            self.reduce_cond = None

        # self.film = AVAILABLE_BLOCKS['film'](512, 128)
        self.film_mul = nn.Linear(
            512 if reduce_cond is None else reduce_cond, reduce_dim
        )
        self.film_add = nn.Linear(
            512 if reduce_cond is None else reduce_cond, reduce_dim
        )

        # DEPRECATED
        # self.conditional_map = {c['id']: c['synonyms'] for c in json.load(open(cond_map))}

        assert len(self.extract_layers) == depth

        self.reduces = nn.ModuleList([nn.Linear(768, reduce_dim) for _ in range(depth)])
        self.blocks = nn.ModuleList(
            [
                nn.TransformerEncoderLayer(d_model=reduce_dim, nhead=n_heads)
                for _ in range(len(self.extract_layers))
            ]
        )
        self.extra_blocks = nn.ModuleList(
            [
                nn.TransformerEncoderLayer(d_model=reduce_dim, nhead=n_heads)
                for _ in range(extra_blocks)
            ]
        )

        trans_conv_ks = (16, 16)
        self.trans_conv = nn.ConvTranspose2d(
            reduce_dim, 1, trans_conv_ks, stride=trans_conv_ks
        )

        # refinement and trans conv

        if learn_trans_conv_only:
            for p in self.parameters():
                p.requires_grad_(False)

            for p in self.trans_conv.parameters():
                p.requires_grad_(True)

        if prompt == "fixed":
            self.prompt_list = ["a photo of a {}."]
        elif prompt == "shuffle":
            self.prompt_list = [
                "a photo of a {}.",
                "a photograph of a {}.",
                "an image of a {}.",
                "{}.",
            ]
        elif prompt == "shuffle+":
            self.prompt_list = [
                "a photo of a {}.",
                "a photograph of a {}.",
                "an image of a {}.",
                "{}.",
                "a cropped photo of a {}.",
                "a good photo of a {}.",
                "a photo of one {}.",
                "a bad photo of a {}.",
                "a photo of the {}.",
            ]
        elif prompt == "shuffle_clip":
            from models.clip_prompts import imagenet_templates

            self.prompt_list = imagenet_templates

        if process_cond is not None:
            if process_cond == "clamp" or process_cond[0] == "clamp":
                val = process_cond[1] if type(process_cond) in {list, tuple} else 0.2

                def clamp_vec(x):
                    return torch.clamp(x, -val, val)

                self.process_cond = clamp_vec

            elif process_cond.endswith(".pth"):
                shift = torch.load(process_cond)

                def add_shift(x):
                    return x + shift.to(x.device)

                self.process_cond = add_shift

        import pickle

        precomp = pickle.load(open("precomputed_prompt_vectors.pickle", "rb"))
        self.precomputed_prompts = {k: torch.from_numpy(v) for k, v in precomp.items()}

    def forward(self, inp_image, conditional=None, return_features=False, mask=None):
        assert type(return_features) == bool

        # inp_image = inp_image.to(self.model.positional_embedding.device)

        if mask is not None:
            raise ValueError("mask not supported")

        # x_inp = normalize(inp_image)
        x_inp = inp_image

        bs, dev = inp_image.shape[0], x_inp.device

        inp_image_size = inp_image.shape[2:]

        cond = self.get_cond_vec(conditional, bs)

        visual_q, activations, _ = self.visual_forward(
            x_inp, extract_layers=[0] + list(self.extract_layers)
        )

        activation1 = activations[0]
        activations = activations[1:]

        a = None
        for i, (activation, block, reduce) in enumerate(
            zip(activations[::-1], self.blocks, self.reduces)
        ):
            if a is not None:
                a = reduce(activation) + a
            else:
                a = reduce(activation)

            if i == self.cond_layer:
                if self.reduce_cond is not None:
                    cond = self.reduce_cond(cond)

                a = self.film_mul(cond) * a + self.film_add(cond)

            a = block(a)

        for block in self.extra_blocks:
            a = a + block(a)

        a = a[1:].permute(1, 2, 0)  # rm cls token and -> BS, Feats, Tokens

        size = int(math.sqrt(a.shape[2]))

        a = a.view(bs, a.shape[1], size, size)

        if self.trans_conv is not None:
            a = self.trans_conv(a)

        if self.upsample_proj is not None:
            a = self.upsample_proj(a)
            a = nnf.interpolate(a, x_inp.shape[2:], mode="bilinear")

        a = nnf.interpolate(a, inp_image_size)

        if return_features:
            return a, visual_q, cond, [activation1] + activations
        else:
            return (a,)