Dynamatrix/DiffBIR-OpenXLab

1import torch2import torch.nn as nn3from torch.utils.checkpoint import checkpoint4 5from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel6 7import open_clip8from ldm.util import default, count_params9 10 11class AbstractEncoder(nn.Module):12    def __init__(self):13        super().__init__()14 15    def encode(self, *args, **kwargs):16        raise NotImplementedError17 18 19class IdentityEncoder(AbstractEncoder):20 21    def encode(self, x):22        return x23 24 25class ClassEmbedder(nn.Module):26    def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1):27        super().__init__()28        self.key = key29        self.embedding = nn.Embedding(n_classes, embed_dim)30        self.n_classes = n_classes31        self.ucg_rate = ucg_rate32 33    def forward(self, batch, key=None, disable_dropout=False):34        if key is None:35            key = self.key36        # this is for use in crossattn37        c = batch[key][:, None]38        if self.ucg_rate > 0. and not disable_dropout:39            mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)40            c = mask * c + (1-mask) * torch.ones_like(c)*(self.n_classes-1)41            c = c.long()42        c = self.embedding(c)43        return c44 45    def get_unconditional_conditioning(self, bs, device="cuda"):46        uc_class = self.n_classes - 1  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)47        uc = torch.ones((bs,), device=device) * uc_class48        uc = {self.key: uc}49        return uc50 51 52def disabled_train(self, mode=True):53    """Overwrite model.train with this function to make sure train/eval mode54    does not change anymore."""55    return self56 57 58class FrozenT5Embedder(AbstractEncoder):59    """Uses the T5 transformer encoder for text"""60    def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77, freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl61        super().__init__()62        self.tokenizer = T5Tokenizer.from_pretrained(version)63        self.transformer = T5EncoderModel.from_pretrained(version)64        self.device = device65        self.max_length = max_length   # TODO: typical value?66        if freeze:67            self.freeze()68 69    def freeze(self):70        self.transformer = self.transformer.eval()71        #self.train = disabled_train72        for param in self.parameters():73            param.requires_grad = False74 75    def forward(self, text):76        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,77                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")78        tokens = batch_encoding["input_ids"].to(self.device)79        outputs = self.transformer(input_ids=tokens)80 81        z = outputs.last_hidden_state82        return z83 84    def encode(self, text):85        return self(text)86 87 88class FrozenCLIPEmbedder(AbstractEncoder):89    """Uses the CLIP transformer encoder for text (from huggingface)"""90    LAYERS = [91        "last",92        "pooled",93        "hidden"94    ]95    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,96                 freeze=True, layer="last", layer_idx=None):  # clip-vit-base-patch3297        super().__init__()98        assert layer in self.LAYERS99        self.tokenizer = CLIPTokenizer.from_pretrained(version)100        self.transformer = CLIPTextModel.from_pretrained(version)101        self.device = device102        self.max_length = max_length103        if freeze:104            self.freeze()105        self.layer = layer106        self.layer_idx = layer_idx107        if layer == "hidden":108            assert layer_idx is not None109            assert 0 <= abs(layer_idx) <= 12110 111    def freeze(self):112        self.transformer = self.transformer.eval()113        #self.train = disabled_train114        for param in self.parameters():115            param.requires_grad = False116 117    def forward(self, text):118        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,119                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")120        tokens = batch_encoding["input_ids"].to(self.device)121        outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer=="hidden")122        if self.layer == "last":123            z = outputs.last_hidden_state124        elif self.layer == "pooled":125            z = outputs.pooler_output[:, None, :]126        else:127            z = outputs.hidden_states[self.layer_idx]128        return z129 130    def encode(self, text):131        return self(text)132 133 134class FrozenOpenCLIPEmbedder(AbstractEncoder):135    """136    Uses the OpenCLIP transformer encoder for text137    """138    LAYERS = [139        #"pooled",140        "last",141        "penultimate"142    ]143    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,144                 freeze=True, layer="last"):145        super().__init__()146        assert layer in self.LAYERS147        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)148        del model.visual149        self.model = model150 151        self.device = device152        self.max_length = max_length153        if freeze:154            self.freeze()155        self.layer = layer156        if self.layer == "last":157            self.layer_idx = 0158        elif self.layer == "penultimate":159            self.layer_idx = 1160        else:161            raise NotImplementedError()162 163    def freeze(self):164        self.model = self.model.eval()165        for param in self.parameters():166            param.requires_grad = False167 168    def forward(self, text):169        tokens = open_clip.tokenize(text)170        z = self.encode_with_transformer(tokens.to(self.device))171        return z172 173    def encode_with_transformer(self, text):174        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]175        x = x + self.model.positional_embedding176        x = x.permute(1, 0, 2)  # NLD -> LND177        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)178        x = x.permute(1, 0, 2)  # LND -> NLD179        x = self.model.ln_final(x)180        return x181 182    def text_transformer_forward(self, x: torch.Tensor, attn_mask = None):183        for i, r in enumerate(self.model.transformer.resblocks):184            if i == len(self.model.transformer.resblocks) - self.layer_idx:185                break186            if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():187                x = checkpoint(r, x, attn_mask)188            else:189                x = r(x, attn_mask=attn_mask)190        return x191 192    def encode(self, text):193        return self(text)194 195 196class FrozenCLIPT5Encoder(AbstractEncoder):197    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda",198                 clip_max_length=77, t5_max_length=77):199        super().__init__()200        self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)201        self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length)202        print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder)*1.e-6:.2f} M parameters, "203              f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder)*1.e-6:.2f} M params.")204 205    def encode(self, text):206        return self(text)207 208    def forward(self, text):209        clip_z = self.clip_encoder.encode(text)210        t5_z = self.t5_encoder.encode(text)211        return [clip_z, t5_z]212 213 214