fix branch

Author: Lupin1998

taming/modules/diffusionmodules/model_mergevq.py

@@ -300,13 +300,15 @@ def apply_merge(self, model, merge_ratio=None, merge_num=None, cand_distribution
             merge_ratio = tome.check_parse_r(num_layers, merge_num, self.embed_res ** 2, inflect)
             if isinstance(cand_distribution, str):
                 cand_distribution = cand_distribution.split('-')
-                if len(cand_distribution) == 2:
-                    cand_distribution, r_cand_num = cand_distribution[0], int(cand_distribution[-1])
+                if len(cand_distribution) == 3:
+                    cand_distribution, r_cand_num, bias = cand_distribution[0], int(cand_distribution[ 1]), int(cand_distribution[-1])
+                elif len(cand_distribution) == 2:
+                    cand_distribution, r_cand_num, bias = cand_distribution[0], int(cand_distribution[-1]), 1
                 else:
-                    cand_distribution, r_cand_num = cand_distribution[0], 5
+                    cand_distribution, r_cand_num, bias = cand_distribution[0], 5, 1
             if cand_distribution is not None:
                 # generate candidate list with the center index
-                bias = 1 if cand_distribution.lower() == "gaussian" else 0
+                bias = bias if cand_distribution.lower() == "gaussian" else 0
                 remain_list = [int(((self.embed_res**2 - merge_num) **0.5 + i) **2) \
                                for i in range(int(-bias), r_cand_num - bias)]
                 merged_list = [self.embed_res**2 - num for num in remain_list]
@@ -359,27 +361,51 @@ def load_pretrained_vit(self, model_path=None):
         weight_selection = {}
         for key in student_weights.keys():
             if ('block' in key or 'cls_token' in key) and key in teacher_weights.keys():
-                weight_selection[key] = uniform_element_selection(teacher_weights[key], student_weights[key].shape)
+                weight_selection[key] = uniform_element_selection(teacher_weights[key], student_weights[key])
         # load to attention
         print("load pre-trained model for encoder:\n",
               self.attn.load_state_dict(weight_selection, strict=False))
 
 
-def uniform_element_selection(tea_weights, stu_shape):
-    """Large Model Initialization (https://arxiv.org/abs/2311.18823)"""
-    assert tea_weights.dim() == len(stu_shape), "Tensors have different number of dimensions"
+def uniform_element_selection(tea_weights, stu_weights):
+    """Modified and borrowed from `Large Model Initialization` (https://arxiv.org/abs/2311.18823)"""
+    assert tea_weights.dim() == stu_weights.dim(), "Tensors have different number of dimensions"
     tea_weights = tea_weights.clone()
-    if tea_weights.shape != stu_shape:
+
+    def interpolate_1d(x, dim, size, up_mode='nearest'):
+        if x.shape[dim] == size:
+            return x
+        permute_order = list(range(x.dim()))
+        permute_order[dim] = x.dim()-1
+        permute_order[-1] = dim
+        x = x.permute(permute_order).contiguous()
+        input_shape = x.shape
+        x = x.view(-1, input_shape[-1]).unsqueeze(1)
+        # upsampling
+        x = torch.nn.functional.interpolate(x, size=size, mode=up_mode) if up_mode == 'nearest' else \
+            torch.nn.functional.interpolate(x, size=size, mode=up_mode, align_corners=False)
+        # reshape back
+        x = x.squeeze(1).view(*input_shape[:-1], size)
+        inv_order = [0] * len(permute_order)
+        for i, o in enumerate(permute_order):
+            inv_order[o] = i
+        return x.permute(inv_order).contiguous()
+
+    if tea_weights.shape != stu_weights.shape:
         for dim in range(tea_weights.dim()):
-            assert tea_weights.shape[dim] >= stu_shape[dim], "Teacher's dimension should not be smaller than students'"
-            if tea_weights.shape[dim] % stu_shape[dim] == 0:
-                step = tea_weights.shape[dim] // stu_shape[dim]
-                indices = torch.arange(stu_shape[dim]) * step
+            if tea_weights.shape[dim] >= stu_weights.shape[dim]:
+                # Teacher's dimension >= students' dimensions
+                if tea_weights.shape[dim] % stu_weights.shape[dim] == 0:
+                    step = tea_weights.shape[dim] // stu_weights.shape[dim]
+                    indices = torch.arange(stu_weights.shape[dim]) * step
+                else:
+                    indices = torch.round(torch.linspace(0, tea_weights.shape[dim]-1, stu_weights.shape[dim])).long()
+                tea_weights = torch.index_select(tea_weights, dim, indices)
             else:
-                indices = torch.round(torch.linspace(0, tea_weights.shape[dim]-1, stu_shape[dim])).long()
-            tea_weights = torch.index_select(tea_weights, dim, indices)
+                # Teacher's dimension < students' dimensions
+                tea_weights = interpolate_1d(tea_weights, dim, stu_weights.shape[dim])
     else:
-        assert tea_weights.shape == stu_shape, "Selected weight should be the same as student"
+        assert tea_weights.shape == stu_weights.shape, "Selected weight should be the same as student"
     return tea_weights
 
 
@@ -732,7 +758,7 @@ def forward(self, x, quantizer):
     from fvcore.nn import FlopCountAnalysis, flop_count_table
     resolution = 256
     cand_distribution = 'gaussian-6'
-    cand_sample_times = 0
+    cand_sample_times = 10
 
     # ch, ch_mult = 64, (1, 2, 4, 8)
     # num_att_blocks, num_res_blocks, r, merge_num = 12, 4, None, 768
@@ -773,6 +799,7 @@ def forward(self, x, quantizer):
     print('encoder (r={}): {}'.format(model.attn.r, y.shape))
     print(flop_count_table(flop, max_depth=4))
     print('MACs (G) of Encoder: {:.3f}'.format(flop.total() / 1e9))
+    # print(model)
 
     if source is not None:
         print('encoder source matrix:', source.shape)

taming/modules/losses/__init__.py

@@ -0,0 +1 @@
+from taming.modules.losses.vqperceptual import DummyLoss

taming/modules/losses/diff_aug.py

@@ -0,0 +1,83 @@
+# Differentiable Augmentation for Data-Efficient GAN Training
+# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han
+# https://arxiv.org/pdf/2006.10738
+
+import torch
+import torch.nn.functional as F
+
+from functools import partial
+
+
+def DiffAugment(x, policy='', prob=0.5, channels_first=True):
+    if policy:
+        if not channels_first:
+            x = x.permute(0, 3, 1, 2)
+        for p in policy.split(','):
+            if torch.rand(1) > prob:
+                continue
+            
+            for f in AUGMENT_FNS[p]:
+                x = f(x)
+            
+        if not channels_first:
+            x = x.permute(0, 2, 3, 1)
+        x = x.contiguous()
+    return x
+
+
+def rand_brightness(x):
+    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+    return x
+
+
+def rand_saturation(x):
+    x_mean = x.mean(dim=1, keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
+    return x
+
+
+def rand_contrast(x):
+    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
+    return x
+
+
+def rand_translation(x, ratio=0.125):
+    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
+    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(x.size(2), dtype=torch.long, device=x.device),
+        torch.arange(x.size(3), dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
+    return x
+
+
+def rand_cutout(x, ratio=0.2):
+    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
+    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
+    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
+    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+    mask[grid_batch, grid_x, grid_y] = 0
+    x = x * mask.unsqueeze(1)
+    return x
+
+
+AUGMENT_FNS = {
+    'color': [rand_brightness, rand_saturation, rand_contrast],
+    'translation': [rand_translation],
+    'cutout_0.2': [partial(rand_cutout, ratio=0.2)],
+    'cutout_0.5': [partial(rand_cutout, ratio=0.5)],
+}

taming/modules/losses/vqperceptual.py

@@ -99,7 +99,6 @@ def __init__(self, disc_start, disc_loss="hinge", disc_dim=64, disc_type="patchg
         super().__init__()
         # discriminator loss
         assert disc_loss in ["hinge", "vanilla", "non_saturate"]
-        assert disc_type == "patchgan"
         self.discriminator = NLayerDiscriminator(
             input_nc=disc_in_channels,
             n_layers=disc_num_layers,