MeshAnything/miche/michelangelo/models/conditional_encoders/encoder_factory.py

# -*- coding: utf-8 -*-
import os

import torch
import torch.nn as nn
from torchvision import transforms
from transformers import CLIPModel, CLIPTokenizer
from collections import OrderedDict

from MeshAnything.miche.michelangelo.data.transforms import RandomResize


class AbstractEncoder(nn.Module):
    embedding_dim: int

    def __init__(self):
        super().__init__()

    def encode(self, *args, **kwargs):
        raise NotImplementedError


class ClassEmbedder(nn.Module):
    def __init__(self, embed_dim, n_classes=1000, key="class"):
        super().__init__()
        self.key = key
        self.embedding = nn.Embedding(n_classes, embed_dim)

    def forward(self, batch, key=None):
        if key is None:
            key = self.key
        # this is for use in crossattn
        c = batch[key][:, None]
        c = self.embedding(c)
        return c


class FrozenCLIPTextEmbedder(AbstractEncoder):
    """Uses the CLIP transformer encoder for text (from Hugging Face)"""

    def __init__(
        self,
        version="openai/clip-vit-large-patch14",
        tokenizer_version=None,
        device="cuda",
        max_length=77,
        zero_embedding_radio: float = 0.1,
    ):
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_version or version)

        self.device = device
        self.max_length = max_length
        self.zero_embedding_radio = zero_embedding_radio

        self.clip_dict = OrderedDict()
        self.clip_name = os.path.split(version)[-1]

        transformer = CLIPModel.from_pretrained(version).text_model

        for param in transformer.parameters():
            param.requires_grad = False
        self.clip_dict[self.clip_name] = transformer

        self._move_flag = False

    @property
    def clip(self):
        return self.clip_dict[self.clip_name]

    def move(self):
        if self._move_flag:
            return

        self.clip_dict[self.clip_name] = self.clip_dict[self.clip_name].to(self.device)
        self._move_flag = True

    def unconditional_embedding(self, batch_size):
        empty_text = [""] * batch_size
        empty_z = self.forward(empty_text)
        return empty_z

    def forward(self, text):
        self.move()

        batch_encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            return_length=True,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )

        tokens = batch_encoding["input_ids"].to(self.device)
        outputs = self.clip(input_ids=tokens)

        z = outputs.last_hidden_state
        return z

    def encode(self, text):
        batch_size = len(text)
        batch_mask = torch.rand((batch_size,))
        for i in range(batch_size):
            if batch_mask[i] < self.zero_embedding_radio:
                text[i] = ""

        return self(text)

class FrozenAlignedCLIPTextEmbedder(AbstractEncoder):
    """Uses the CLIP transformer encoder for text (from Hugging Face)"""

    def __init__(
        self,
        version="openai/clip-vit-large-patch14",
        tokenizer_version=None,
        device="cuda",
        max_length=77,
        zero_embedding_radio: float = 0.1,
    ):
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_version or version)

        self.device = device
        self.max_length = max_length
        self.zero_embedding_radio = zero_embedding_radio

        self.clip_dict = OrderedDict()
        self.clip_name = os.path.split(version)[-1]

        transformer = CLIPModel.from_pretrained(version).text_model

        for param in transformer.parameters():
            param.requires_grad = False
        self.clip_dict[self.clip_name] = transformer

        self._move_flag = False

    @property
    def clip(self):
        return self.clip_dict[self.clip_name]

    def move(self):
        if self._move_flag:
            return

        self.clip_dict[self.clip_name] = self.clip_dict[self.clip_name].to(self.device)
        self._move_flag = True

    def unconditional_embedding(self, batch_size):
        empty_text = [""] * batch_size
        empty_z = self.forward(empty_text)
        return empty_z

    def forward(self, text):
        self.move()

        batch_encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            return_length=True,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )

        tokens = batch_encoding["input_ids"].to(self.device)
        outputs = self.clip(input_ids=tokens)

        z = outputs.last_hidden_state
        return z

    def encode(self, text):
        batch_size = len(text)
        batch_mask = torch.rand((batch_size,))
        for i in range(batch_size):
            if batch_mask[i] < self.zero_embedding_radio:
                text[i] = ""

        return self(text)


class FrozenCLIPImageEmbedder(AbstractEncoder):
    """Uses the CLIP transformer encoder for text (from Hugging Face)"""

    def __init__(
            self,
            version="openai/clip-vit-large-patch14",
            device="cuda",
            zero_embedding_radio=0.1,
            normalize_embedding=True,
            num_projection_vector=0,
            linear_mapping_bias=True,
            reverse_visual_projection=False,
    ):
        super().__init__()

        self.device = device

        self.clip_dict = OrderedDict()
        self.clip_name = os.path.split(version)[-1]

        clip_model = CLIPModel.from_pretrained(version)
        clip_model.text_model = None
        clip_model.text_projection = None
        clip_model = clip_model.eval()
        for param in self.parameters():
            param.requires_grad = False
        self.clip_dict[self.clip_name] = clip_model

        self.transform = transforms.Compose(
            [
                transforms.Resize(224, transforms.InterpolationMode.BICUBIC, antialias=True),
                transforms.CenterCrop(224),  # crop a (224, 224) square
                transforms.Normalize(
                    mean=[0.48145466, 0.4578275, 0.40821073],
                    std=[0.26862954, 0.26130258, 0.27577711],
                ),
            ]
        )
        self.zero_embedding_radio = zero_embedding_radio

        self.num_projection_vector = num_projection_vector
        self.reverse_visual_projection = reverse_visual_projection
        self.normalize_embedding = normalize_embedding

        embedding_dim = (
            clip_model.visual_projection.in_features
            if reverse_visual_projection
            else clip_model.visual_projection.out_features
        )
        self.embedding_dim = embedding_dim
        if self.num_projection_vector > 0:
            self.projection = nn.Linear(
                embedding_dim,
                clip_model.visual_projection.out_features * num_projection_vector,
                bias=linear_mapping_bias,
            )
            nn.init.normal_(self.projection.weight, std=embedding_dim ** -0.5)

        self._move_flag = False

    @property
    def clip(self):
        return self.clip_dict[self.clip_name]

    def unconditional_embedding(self, batch_size):
        zero = torch.zeros(
            batch_size,
            1,
            self.embedding_dim,
            device=self.device,
            dtype=self.clip.visual_projection.weight.dtype,
        )
        if self.num_projection_vector > 0:
            zero = self.projection(zero).view(batch_size, self.num_projection_vector, -1)
        return zero

    def forward(self, image, value_range=(-1, 1), zero_embedding_radio=0):
        if value_range is not None:
            low, high = value_range
            image = (image - low) / (high - low)

        image = image.to(self.device, dtype=self.clip.visual_projection.weight.dtype)

        if self.reverse_visual_projection:
            z = self.clip.vision_model(self.transform(image))[1]
        else:
            z = self.clip.get_image_features(self.transform(image))

        if self.normalize_embedding:
            z = z / z.norm(dim=-1, keepdim=True)
        if z.ndim == 2:
            z = z.unsqueeze(dim=-2)

        if zero_embedding_radio > 0:
            mask = torch.rand((len(image), 1, 1), device=z.device, dtype=z.dtype) < zero_embedding_radio
            z = z * mask.to(z)

        if self.num_projection_vector > 0:
            z = self.projection(z).view(len(image), self.num_projection_vector, -1)

        return z

    def move(self):
        if self._move_flag:
            return

        self.clip_dict[self.clip_name] = self.clip_dict[self.clip_name].to(self.device)
        self._move_flag = True

    def encode(self, image):
        self.move()
        return self(image, zero_embedding_radio=self.zero_embedding_radio)


class FrozenCLIPImageGridEmbedder(AbstractEncoder):

    def __init__(
            self,
            version="openai/clip-vit-large-patch14",
            device="cuda",
            zero_embedding_radio=0.1,
    ):
        super().__init__()

        self.device = device

        self.clip_dict = OrderedDict()
        self.clip_name = os.path.split(version)[-1]

        clip_model: CLIPModel = CLIPModel.from_pretrained(version)
        clip_model.text_model = None
        clip_model.text_projection = None
        clip_model = clip_model.eval()
        for param in self.parameters():
            param.requires_grad = False
        self.clip_dict[self.clip_name] = clip_model

        self.transform = transforms.Compose(
            [
                transforms.Resize(224, transforms.InterpolationMode.BILINEAR, antialias=True),
                transforms.CenterCrop(224),  # crop a (224, 224) square
                transforms.Normalize(
                    mean=[0.48145466, 0.4578275, 0.40821073],
                    std=[0.26862954, 0.26130258, 0.27577711],
                ),
            ]
        )
        self.zero_embedding_radio = zero_embedding_radio
        self.embedding_dim = clip_model.vision_embed_dim

        self._move_flag = False

    @property
    def clip(self):
        return self.clip_dict[self.clip_name]

    def move(self):
        if self._move_flag:
            return

        self.clip_dict[self.clip_name] = self.clip_dict[self.clip_name].to(self.device)
        self._move_flag = True

    def unconditional_embedding(self, batch_size):
        zero = torch.zeros(
            batch_size,
            self.clip.vision_model.embeddings.num_positions,
            self.embedding_dim,
            device=self.device,
            dtype=self.clip.visual_projection.weight.dtype,
        )
        return zero

    def forward(self, image, value_range=(-1, 1), zero_embedding_radio=0):
        self.move()

        if value_range is not None:
            low, high = value_range
            image = (image - low) / (high - low)

        image = image.to(self.device, dtype=self.clip.visual_projection.weight.dtype)

        z = self.clip.vision_model(self.transform(image)).last_hidden_state

        if zero_embedding_radio > 0:
            mask = torch.rand((len(image), 1, 1), device=z.device, dtype=z.dtype) >= zero_embedding_radio
            z = z * mask.to(z)

        return z

    def encode(self, image):
        return self(image, zero_embedding_radio=self.zero_embedding_radio)


class MoECLIPImageEncoder(nn.Module):
    def __init__(
            self,
            versions,
            hidden_state_dim,
            num_projection_vector=8,
            zero_embedding_radio=0.1,
            device="cuda",
            precision="fp16",
            normalize=False,
            clip_max=0,
            transform_type="base",
            argument_p=0.2,
    ):
        super().__init__()

        self.device = torch.device(device)
        self.hidden_state_dim = hidden_state_dim
        self.zero_embedding_radio = zero_embedding_radio
        self.num_projection_vector = num_projection_vector
        self.dtype = dict(fp16=torch.float16, fp32=torch.float32, bf16=torch.bfloat16)[precision]
        self.normalize = normalize
        self.clip_max = clip_max

        if transform_type == "base":
            self.transform = transforms.Compose(
                [
                    transforms.Resize(224, transforms.InterpolationMode.BICUBIC, antialias=True),
                    transforms.CenterCrop(224),  # crop a (224, 224) square
                    transforms.Normalize(
                        mean=[0.48145466, 0.4578275, 0.40821073],
                        std=[0.26862954, 0.26130258, 0.27577711],
                    ),
                ]
            )
        elif transform_type == "crop_blur_resize":
            self.transform = transforms.Compose(
                [
                    transforms.Resize(224, transforms.InterpolationMode.BICUBIC, antialias=True),
                    transforms.CenterCrop(224),  # crop a (224, 224) square
                    transforms.RandomApply(
                        transforms=[
                            transforms.RandomResizedCrop(
                                size=224,
                                scale=(0.8, 1.0),
                                ratio=(0.99, 1.01),
                                interpolation=transforms.InterpolationMode.BICUBIC,
                            ),
                        ],
                        p=argument_p,
                    ),
                    transforms.RandomApply(
                        transforms=[
                            transforms.GaussianBlur(kernel_size=9, sigma=(0.1, 5)),
                        ],
                        p=argument_p,
                    ),
                    transforms.RandomApply(
                        transforms=[
                            RandomResize(size=224, resize_radio=(0.2, 1)),
                        ],
                        p=argument_p,
                    ),
                    transforms.Normalize(
                        mean=[0.48145466, 0.4578275, 0.40821073],
                        std=[0.26862954, 0.26130258, 0.27577711],
                    ),
                ]
            )
        else:
            raise ValueError(f"invalid {transform_type=}")

        if isinstance(versions, str):
            versions = (versions,)

        # 如果直接把clips定位为当前类的子module，1. 会在保存ckp时存无用的多个权重。 2. pl会调用to，导致layer_norm的权重也被转换成fp16
        clips = OrderedDict()

        for v in versions:
            # 因为clips不是子module，直接指定device="cuda"会错误地导致clip模型权重都被放到cuda:0上。
            clips[v], _ = clip.load(name=v, device="cpu", jit=False, download_root=None)
            delattr(clips[v], "transformer")
            clips[v].eval()
            clips[v].requires_grad_(False)

        self.clips_hidden_dim = sum(clips[v].ln_final.weight.size(0) for v in clips)

        if self.num_projection_vector == 0:
            self.projection = nn.Identity()
        else:
            self.projection = nn.Linear(self.clips_hidden_dim, hidden_state_dim * self.num_projection_vector, bias=True)
            self.projection.to(dtype=self.dtype)
            nn.init.normal_(self.projection.weight, std=self.clips_hidden_dim ** -0.5)

        self.clips = clips

        self._move_flag = False

    def move(self):
        if self._move_flag:
            return

        def convert_weights(model: nn.Module):
            """Convert applicable model parameters to fp16"""

            def _convert_weights_to_fp16(l):
                if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
                    l.weight.data = l.weight.data.type(self.dtype)
                    if l.bias is not None:
                        l.bias.data = l.bias.data.type(self.dtype)

                if isinstance(l, nn.MultiheadAttention):
                    for attr in [
                        *[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
                        "in_proj_bias",
                        "bias_k",
                        "bias_v",
                    ]:
                        tensor = getattr(l, attr)
                        if tensor is not None:
                            tensor.data = tensor.data.type(self.dtype)

                for name in ["text_projection", "proj"]:
                    if hasattr(l, name):
                        attr = getattr(l, name)
                        if attr is not None:
                            attr.data = attr.data.type(self.dtype)

            model.apply(_convert_weights_to_fp16)

        for k in self.clips:
            self.clips[k].to(self.device)
            convert_weights(self.clips[k])  # fp32 -> self.dtype
        self._move_flag = True

    def unconditional_embedding(self, batch_size=None):
        zero = torch.zeros(
            batch_size,
            self.clips_hidden_dim,
            device=self.device,
            dtype=self.dtype,
        )
        if self.num_projection_vector > 0:
            zero = self.projection(zero).view(batch_size, self.num_projection_vector, -1)
        return zero

    def convert_embedding(self, z):
        if self.num_projection_vector > 0:
            z = self.projection(z.type(self.projection.weight.dtype)).view(len(z), self.num_projection_vector, -1)
        return z

    def forward(self, image, value_range=(-1, 1), zero_embedding_radio=0):
        if value_range is not None:
            low, high = value_range
            image = (image - low) / (high - low)

        image = self.transform(image)

        with torch.no_grad():
            embs = []
            for v in self.clips:
                x = self.clips[v].encode_image(image)
                if self.normalize:
                    x = x / x.norm(p=2, dim=-1, keepdim=True) * (x.size(-1) ** 0.5)
                    # clip_max only works with normalization
                    if self.clip_max > 0:
                        x = x.clamp(-self.clip_max, self.clip_max)
                embs.append(x)

            z = torch.cat(embs, dim=-1)
            if self.normalize:
                z /= z.size(-1) ** 0.5

        if zero_embedding_radio > 0:
            mask = torch.rand((len(image), 1, 1), device=z.device, dtype=z.dtype) >= zero_embedding_radio
            z = z + mask.to(z)

        if self.num_projection_vector > 0:
            z = self.projection(z).view(len(image), self.num_projection_vector, -1)
        return z

    def encode(self, image):
        self.move()
        return self(image, zero_embedding_radio=self.zero_embedding_radio)