jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x, position_embedding=None):
        x = self.proj(x)

        if position_embedding is not None:
            # interpolate the position embedding to the corresponding size
            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3, 1, 2)
            Hp, Wp = x.shape[2], x.shape[3]
            position_embedding = F.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
            x = x + position_embedding

        x = x.flatten(2).transpose(1, 2)
        return x
//  [python]
def forward(self, x):
        if self.training:
            self.pca_computed = False
            if self.downsampling is not None:
                x1 = t(x, size=self.downsampling, mode='nearest')
            else:
                x1 = x
            swapped: torch.Tensor = x1.permute([1, 0, 2, 3])
            flattened: torch.Tensor = swapped.flatten(1)
            reshaped_batch: torch.Tensor = flattened.permute([1, 0])
            self._update_autorcorrelation(reshaped_batch)
            return x
        else:
            if not self.pca_computed:
                self._compute_autorcorrelation()
                self._compute_eigenspace()
                self._compute_pca_matrix()
                self._reset_autorcorrelation()
                self.pca_computed = True
            x1 = self.mean_subtracting_convolution(x)
            x = x + x1
            return self.convolution(x)
//  [python]
def forward(self, x):
        if self.training:
            self.pca_computed = False
            if self.downsampling is not None:
                x1 = interpolate(x, size=self.downsampling, mode='nearest')
            else:
                x1 = x
            swapped: torch.Tensor = x1.permute([1, 0, 2, 3])
            flattened: torch.Tensor = swapped.flatten(1)
            reshaped_batch: torch.Tensor = flattened.permute([1, 0])
            self._update_autorcorrelation(reshaped_batch)
            return x
        else:
            if not self.pca_computed:
                self._compute_autorcorrelation()
                self._compute_eigenspace()
                self._compute_pca_matrix()
                self._reset_autorcorrelation()
                self.pca_computed = True
            x1 = self.mean_subtracting_convolution(x)
            x = x + x1
            return self.convolution(x)
//  [python]
def forward(self, xs):
        """Forward pass through CBFuse layer."""
        target_size = xs[-1].shape[2:]
        res = [t.interpolate(x[self.idx[i]], size=target_size, mode="nearest") for i, x in enumerate(xs[:-1])]
        out = torch.sum(torch.stack(res + xs[-1:]), dim=0)
        return out
//  [python]
def forward(self, xs):
        """Forward pass through CBFuse layer."""
        target_size = xs[-1].shape[2:]
        res = [F.interpolate(x[self.idx[i]], size=target_size, mode="nearest") for i, x in enumerate(xs[:-1])]
        out = torch.sum(torch.stack(res + xs[-1:]), dim=0)
        return out
//  [python]
def forward(self, x):
        """ Base network """
        if self.amp:
            with torch.cuda.amp.autocast():
                sources = self.basenet(x)

                """ U network """
                y = torch.cat([sources[0], sources[1]], dim=1)
                y = self.upconv1(y)

                y = F.upsample(y, size=sources[2].size()[2:], mode='bilinear')
                y = torch.cat([y, sources[2]], dim=1)
                y = self.upconv2(y)

                y = F.upsample(y, size=sources[3].size()[2:], mode='bilinear')
                y = torch.cat([y, sources[3]], dim=1)
                y = self.upconv3(y)

                y = F.upsample(y, size=sources[4].size()[2:], mode='bilinear')
                y = torch.cat([y, sources[4]], dim=1)
                feature = self.upconv4(y)

                y = self.conv_cls(feature)

                return y.permute(0,2,3,1), feature
        else:

            sources = self.basenet(x)

            """ U network """
            y = torch.cat([sources[0], sources[1]], dim=1)
            y = self.upconv1(y)

            y = F.upsample(y, size=sources[2].size()[2:], mode='bilinear')
            y = torch.cat([y, sources[2]], dim=1)
            y = self.upconv2(y)

            y = F.upsample(y, size=sources[3].size()[2:], mode='bilinear')
            y = torch.cat([y, sources[3]], dim=1)
            y = self.upconv3(y)

            y = F.upsample(y, size=sources[4].size()[2:], mode='bilinear')
            y = torch.cat([y, sources[4]], dim=1)
            feature = self.upconv4(y)

            y = self.conv_cls(feature)

            return y.permute(0, 2, 3, 1), feature
//  [python]
def forward(self, x):
        """ Base network """
        if self.amp:
            with torch.cuda.amp.autocast():
                sources = self.basenet(x)

                """ U network """
                y = torch.cat([sources[0], sources[1]], dim=1)
                y = self.upconv1(y)

                y = F.interpolate(y, size=sources[2].size()[2:], mode='bilinear', align_corners=False)
                y = torch.cat([y, sources[2]], dim=1)
                y = self.upconv2(y)

                y = F.interpolate(y, size=sources[3].size()[2:], mode='bilinear', align_corners=False)
                y = torch.cat([y, sources[3]], dim=1)
                y = self.upconv3(y)

                y = F.interpolate(y, size=sources[4].size()[2:], mode='bilinear', align_corners=False)
                y = torch.cat([y, sources[4]], dim=1)
                feature = self.upconv4(y)

                y = self.conv_cls(feature)

                return y.permute(0,2,3,1), feature
        else:

            sources = self.basenet(x)

            """ U network """
            y = torch.cat([sources[0], sources[1]], dim=1)
            y = self.upconv1(y)

            y = F.interpolate(y, size=sources[2].size()[2:], mode='bilinear', align_corners=False)
            y = torch.cat([y, sources[2]], dim=1)
            y = self.upconv2(y)

            y = F.interpolate(y, size=sources[3].size()[2:], mode='bilinear', align_corners=False)
            y = torch.cat([y, sources[3]], dim=1)
            y = self.upconv3(y)

            y = F.interpolate(y, size=sources[4].size()[2:], mode='bilinear', align_corners=False)
            y = torch.cat([y, sources[4]], dim=1)
            feature = self.upconv4(y)

            y = self.conv_cls(feature)

            return y.permute(0, 2, 3, 1), feature
//  [python]
def forward(self, x):
        if self.scale_factor == 2.0:
            x = F.upsample(x, scale_factor=2, mode='nearest')
        elif self.scale_factor == 0.5:
            x = F.avg_pool2d(x, kernel_size=2, stride=2)
        return x
//  [python]
def forward(self, x):
        if self.scale_factor == 2.0:
            x = F.interpolate(x, scale_factor=2, mode='nearest')
        elif self.scale_factor == 0.5:
            x = F.avg_pool2d(x, kernel_size=2, stride=2)
        return x
//  [python]
def forward_raw(self, x):
        """Forward function."""
        x = self.patch_embed(x)

        Wh, Ww = x.size(2), x.size(3)
        if self.ape:
            # interpolate the position embedding to the corresponding size
            absolute_pos_embed = F.interpolate(
                self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
            )
            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
        else:
            x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        outs = []
        for i in range(self.num_layers):
            layer = self.layers[i]
            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
            # import ipdb; ipdb.set_trace()

            if i in self.out_indices:
                norm_layer = getattr(self, f"norm{i}")
                x_out = norm_layer(x_out)

                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
                outs.append(out)
        # in:
        #   torch.Size([2, 3, 1024, 1024])
        # outs:
        #   [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \
        #       torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])]
        return tuple(outs)
//  [python]
def invoke(self, context: InvocationContext) -> LatentsOutput:
        latents = context.tensors.load(self.latents.latents_name)
        device = TorchDevice.choose_torch_device()

        resized_latents = torch.nn.functional.interpolate(
            latents.to(device),
            size=(self.height // LATENT_SCALE_FACTOR, self.width // LATENT_SCALE_FACTOR),
            mode=self.mode,
            antialias=self.antialias if self.mode in ["bilinear", "bicubic"] else False,
        )

        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
        resized_latents = resized_latents.to("cpu")

        TorchDevice.empty_cache()

        name = context.tensors.save(tensor=resized_latents)
        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
//  [python]
def invoke(self, context: InvocationContext) -> LatentsOutput:
        latents = context.tensors.load(self.latents.latents_name)

        device = TorchDevice.choose_torch_device()

        # resizing
        resized_latents = t.interpolate(
            latents.to(device),
            scale_factor=self.scale_factor,
            mode=self.mode,
            align_corners=self.antialias if self.mode in ["bilinear", "bicubic"] else False,
        )

        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
        resized_latents = resized_latents.to("cpu")
        TorchDevice.empty_cache()

        name = context.tensors.save(tensor=resized_latents)
        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
//  [python]
def invoke(self, context: InvocationContext) -> LatentsOutput:
        latents = context.tensors.load(self.latents.latents_name)

        device = TorchDevice.choose_torch_device()

        # resizing
        resized_latents = torch.nn.functional.interpolate(
            latents.to(device),
            scale_factor=self.scale_factor,
            mode=self.mode,
            antialias=self.antialias if self.mode in ["bilinear", "bicubic"] else False,
        )

        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
        resized_latents = resized_latents.to("cpu")
        TorchDevice.empty_cache()

        name = context.tensors.save(tensor=resized_latents)
        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
//  [python]
def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors

        """Forward function."""
        x = self.patch_embed(x)

        Wh, Ww = x.size(2), x.size(3)
        if self.ape:
            # interpolate the position embedding to the corresponding size
            absolute_pos_embed = F.interpolate(
                self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
            )
            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
        else:
            x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        outs = []
        for i in range(self.num_layers):
            layer = self.layers[i]
            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)

            if i in self.out_indices:
                norm_layer = getattr(self, f"norm{i}")
                x_out = norm_layer(x_out)

                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
                outs.append(out)
        # in:
        #   torch.Size([2, 3, 1024, 1024])
        # out:
        #   [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \
        #       torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])]

        # collect for nesttensors
        outs_dict = {}
        for idx, out_i in enumerate(outs):
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0]
            outs_dict[idx] = NestedTensor(out_i, mask)

        return outs_dict
//  [python]
def forward(self, feature, image):
        cp = self.compress(feature)
        seg = torch.cat((F.upsample(cp, size=(image.shape[2], image.shape[3])), image), dim=1)   # 7, 256, 256
    
        x = feature      # 256, 64, 64
        x = self.fc(x)                # 512, 64, 64
        x = self.G_middle_0(x, seg)
        x = self.G_middle_1(x, seg)
        x = self.G_middle_2(x, seg)
        # x = self.G_middle_3(x, seg)
        # x = self.G_middle_4(x, seg)
        # x = self.G_middle_5(x, seg)
        x = self.up(x)                # 256, 128, 128
        x = self.up_0(x, seg)
        x = self.up(x)                # 64, 256, 256
        x = self.up_1(x, seg)
        x = self.conv_img(F.leaky_relu(x, 2e-1))
        # x = torch.tanh(x)
        x = F.sigmoid(x)
        
        return x
//  [python]
def forward(self, feature, image):
        cp = self.compress(feature)
        seg = torch.cat((F.interpolate(cp, size=(image.shape[2], image.shape[3])), image), dim=1)   # 7, 256, 256
    
        x = feature      # 256, 64, 64
        x = self.fc(x)                # 512, 64, 64
        x = self.G_middle_0(x, seg)
        x = self.G_middle_1(x, seg)
        x = self.G_middle_2(x, seg)
        # x = self.G_middle_3(x, seg)
        # x = self.G_middle_4(x, seg)
        # x = self.G_middle_5(x, seg)
        x = self.up(x)                # 256, 128, 128
        x = self.up_0(x, seg)
        x = self.up(x)                # 64, 256, 256
        x = self.up_1(x, seg)
        x = self.conv_img(F.leaky_relu(x, 2e-1))
        # x = torch.tanh(x)
        x = F.sigmoid(x)
        
        return x
//  [python]
def checkpoint_filter_fn(state_dict, model):
    if 'stages.0.blocks.0.norm1.weight' in state_dict:
        # native checkpoint, look for rel_pos interpolations
        for k in state_dict.keys():
            if 'rel_pos' in k:
                rel_pos = state_dict[k]
                dest_rel_pos_shape = model.state_dict()[k].shape
                if rel_pos.shape[0] != dest_rel_pos_shape[0]:
                    rel_pos_resized = t.nn.functional.interpolate(
                        rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
                        size=dest_rel_pos_shape[0],
                        mode="linear",
                    )
                    state_dict[k] = rel_pos_resized.reshape(-1, dest_rel_pos_shape[0]).permute(1, 0)
        return state_dict

    import re
    if 'model_state' in state_dict:
        state_dict = state_dict['model_state']

    depths = getattr(model, 'depths', None)
    expand_attn = getattr(model, 'expand_attn', True)
    assert depths is not None, 'model requires depth attribute to remap checkpoints'
    depth_map = {}
    block_idx = 0
    for stage_idx, d in enumerate(depths):
        depth_map.update({i: (stage_idx, i - block_idx) for i in range(block_idx, block_idx + d)})
        block_idx += d

    out_dict = {}
    for k, v in state_dict.items():
        k = re.sub(
            r'blocks\.(\d+)',
            lambda x: f'stages.{depth_map[int(x.group(1))][0]}.blocks.{depth_map[int(x.group(1))][1]}',
            k)

        if expand_attn:
            k = re.sub(r'stages\.(\d+).blocks\.(\d+).proj', f'stages.\\1.blocks.\\2.shortcut_proj_attn', k)
        else:
            k = re.sub(r'stages\.(\d+).blocks\.(\d+).proj', f'stages.\\1.blocks.\\2.shortcut_proj_mlp', k)
        if 'head' in k:
            k = k.replace('head.projection', 'head.fc')
        out_dict[k] = v

    return out_dict
//  [python]
def checkpoint_filter_fn(state_dict, model):
    if 'stages.0.blocks.0.norm1.weight' in state_dict:
        # native checkpoint, look for rel_pos interpolations
        for k in state_dict.keys():
            if 'rel_pos' in k:
                rel_pos = state_dict[k]
                dest_rel_pos_shape = model.state_dict()[k].shape
                if rel_pos.shape[0] != dest_rel_pos_shape[0]:
                    rel_pos_resized = torch.nn.functional.interpolate(
                        rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
                        size=dest_rel_pos_shape[0],
                        mode="linear",
                    )
                    state_dict[k] = rel_pos_resized.reshape(-1, dest_rel_pos_shape[0]).permute(1, 0)
        return state_dict

    import re
    if 'model_state' in state_dict:
        state_dict = state_dict['model_state']

    depths = getattr(model, 'depths', None)
    expand_attn = getattr(model, 'expand_attn', True)
    assert depths is not None, 'model requires depth attribute to remap checkpoints'
    depth_map = {}
    block_idx = 0
    for stage_idx, d in enumerate(depths):
        depth_map.update({i: (stage_idx, i - block_idx) for i in range(block_idx, block_idx + d)})
        block_idx += d

    out_dict = {}
    for k, v in state_dict.items():
        k = re.sub(
            r'blocks\.(\d+)',
            lambda x: f'stages.{depth_map[int(x.group(1))][0]}.blocks.{depth_map[int(x.group(1))][1]}',
            k)

        if expand_attn:
            k = re.sub(r'stages\.(\d+).blocks\.(\d+).proj', f'stages.\\1.blocks.\\2.shortcut_proj_attn', k)
        else:
            k = re.sub(r'stages\.(\d+).blocks\.(\d+).proj', f'stages.\\1.blocks.\\2.shortcut_proj_mlp', k)
        if 'head' in k:
            k = k.replace('head.projection', 'head.fc')
        out_dict[k] = v

    return out_dict
//  [python]
def transform_tensor(self, image_tensor):
    image_tensor = t.interpolate(
        image_tensor, size=(224, 224), mode="bicubic", align_corners=False
    )
    normalize = transforms.Normalize(
        mean=[0.48145466, 0.4578275, 0.40821073],
        std=[0.26862954, 0.26130258, 0.27577711],
    )
    image_tensor = normalize(image_tensor)
    return image_tensor
//  [python]
def transform_tensor(self, image_tensor):
    image_tensor = torch.nn.functional.interpolate(
        image_tensor, size=(224, 224), mode="bicubic", align_corners=False
    )
    normalize = transforms.Normalize(
        mean=[0.48145466, 0.4578275, 0.40821073],
        std=[0.26862954, 0.26130258, 0.27577711],
    )
    image_tensor = normalize(image_tensor)
    return image_tensor
//  [python]
def forward(self, source_image, kp_driving, kp_source, source_depth, driving_depth):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)

        src_out = self.src_first(source_depth)
        for i in range(len(self.src_down_blocks)):
            src_out = self.src_down_blocks[i](src_out)
        
        # dst_out = self.dst_first(driving_depth)
        # for i in range(len(self.down_blocks)):
        #     dst_out = self.dst_down_blocks[i](dst_out)
        
        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(source_image=source_image, kp_driving=kp_driving,
                                                     kp_source=kp_source)
            output_dict['mask'] = dense_motion['mask']
            output_dict['sparse_deformed'] = dense_motion['sparse_deformed']

            if 'occlusion_map' in dense_motion:
                occlusion_map = dense_motion['occlusion_map']
                output_dict['occlusion_map'] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion['deformation']
            out = self.deform_input(out, deformation)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = t.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
                out = out * occlusion_map
            out,attention  = self.AttnModule(src_out,out)

            output_dict["deformed"] = self.deform_input(source_image, deformation)
            output_dict["attention"] = attention

        # Decoding part
        out = self.bottleneck(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def forward(self, source_image, kp_driving, kp_source, source_depth, driving_depth):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)

        src_out = self.src_first(source_depth)
        for i in range(len(self.src_down_blocks)):
            src_out = self.src_down_blocks[i](src_out)
        
        # dst_out = self.dst_first(driving_depth)
        # for i in range(len(self.down_blocks)):
        #     dst_out = self.dst_down_blocks[i](dst_out)
        
        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(source_image=source_image, kp_driving=kp_driving,
                                                     kp_source=kp_source)
            output_dict['mask'] = dense_motion['mask']
            output_dict['sparse_deformed'] = dense_motion['sparse_deformed']

            if 'occlusion_map' in dense_motion:
                occlusion_map = dense_motion['occlusion_map']
                output_dict['occlusion_map'] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion['deformation']
            out = self.deform_input(out, deformation)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
                out = out * occlusion_map
            out,attention  = self.AttnModule(src_out,out)

            output_dict["deformed"] = self.deform_input(source_image, deformation)
            output_dict["attention"] = attention

        # Decoding part
        out = self.bottleneck(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def _upsample_like(src, tar):

    src = t.interpolate(src, size=tar.shape[2:], mode="bilinear", align_corners=False)

    return src
//  [python]
def _upsample_like(src, tar):

    src = F.interpolate(src, size=tar.shape[2:], mode="bilinear", align_corners=False)

    return src
//  [python]
def sds_loss(
        self,
        text_embeddings: Float[Tensor, "N max_length embed_dim"],
        image: Float[Tensor, "BS 3 H W"],
        guidance_scale: float = 100.0,
        grad_scaler: Optional[GradScaler] = None,
    ) -> torch.Tensor:
        """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
        Args:
            text_embeddings: Text embeddings
            image: Rendered image
            guidance_scale: How much to weigh the guidance
            grad_scaler: Grad scaler
        Returns:
            The loss
        """
        image = t(image.half(), size=(IMG_DIM, IMG_DIM), mode="bilinear", align_corners=False)
        t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)

        # predict the noise residual with unet, NO grad!
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(image)
            image_noisy = self.scheduler.add_noise(image, noise, t)  # type: ignore
            # pred noise
            image_model_input = torch.cat((image_noisy,) * 2)
            noise_pred = self.unet(image_model_input, t, encoder_hidden_states=text_embeddings).sample

        # perform guidance
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
        noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)

        noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # w(t), sigma_t^2
        w = 1 - self.alphas[t]

        grad = w * (noise_pred - noise)
        grad = torch.nan_to_num(grad)

        target = (image - grad).detach()
        loss = 0.5 * F.mse_loss(image, target, reduction="sum") / image.shape[0]

        return loss
//  [python]
def sds_loss(
        self,
        text_embeddings: Float[Tensor, "N max_length embed_dim"],
        image: Float[Tensor, "BS 3 H W"],
        guidance_scale: float = 100.0,
        grad_scaler: Optional[GradScaler] = None,
    ) -> torch.Tensor:
        """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
        Args:
            text_embeddings: Text embeddings
            image: Rendered image
            guidance_scale: How much to weigh the guidance
            grad_scaler: Grad scaler
        Returns:
            The loss
        """
        image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear", align_corners=False)
        t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)

        # predict the noise residual with unet, NO grad!
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(image)
            image_noisy = self.scheduler.add_noise(image, noise, t)  # type: ignore
            # pred noise
            image_model_input = torch.cat((image_noisy,) * 2)
            noise_pred = self.unet(image_model_input, t, encoder_hidden_states=text_embeddings).sample

        # perform guidance
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
        noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)

        noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # w(t), sigma_t^2
        w = 1 - self.alphas[t]

        grad = w * (noise_pred - noise)
        grad = torch.nan_to_num(grad)

        target = (image - grad).detach()
        loss = 0.5 * F.mse_loss(image, target, reduction="sum") / image.shape[0]

        return loss
//  [python]
def viz_hidden_features(self, last_img_features, current_imgs_features):
        lif = last_img_features.norm(dim=0)
        cif = current_imgs_features.norm(dim=0)
        lif = lif / lif.max()
        cif = cif / cif.max()
        new_size =(self.img_height, self.img_width)
        lif = t.nn.functional.interpolate(lif[None,None], size=new_size, mode="bilinear", align_corners=True)
        cif = t.nn.functional.interpolate(cif[None,None], size=new_size, mode="bilinear", align_corners=True)
        lif = lif.squeeze()
        cif = cif.squeeze()
        cv2.imshow('Last Img Features', lif.cpu().numpy())
        cv2.imshow('Current Img Features', cif.cpu().numpy())
//  [python]
def viz_hidden_features(self, last_img_features, current_imgs_features):
        lif = last_img_features.norm(dim=0)
        cif = current_imgs_features.norm(dim=0)
        lif = lif / lif.max()
        cif = cif / cif.max()
        new_size =(self.img_height, self.img_width)
        lif = F.interpolate(lif[None,None], size=new_size, mode="bilinear", align_corners=True)
        cif = F.interpolate(cif[None,None], size=new_size, mode="bilinear", align_corners=True)
        lif = lif.squeeze()
        cif = cif.squeeze()
        cv2.imshow('Last Img Features', lif.cpu().numpy())
        cv2.imshow('Current Img Features', cif.cpu().numpy())
//  [python]
def predict(self, x: Tuple[Tensor], batch_img_metas: List[dict],
                test_cfg: ConfigType) -> Tuple[Tensor]:
        """Test without augmentaton.

        Args:
            x (tuple[Tensor]): Multi-level features from the
                upstream network, each is a 4D-tensor.
            batch_img_metas (List[:obj:`SegDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_sem_seg`.
            test_cfg (ConfigType): Test config.

        Returns:
            Tensor: A tensor of segmentation mask.
        """

        batch_data_samples = []
        for metainfo in batch_img_metas:
            metainfo['batch_input_shape'] = metainfo['img_shape']
            batch_data_samples.append(SegDataSample(metainfo=metainfo))
        # Forward function of MaskFormerHead from MMDetection needs
        # 'batch_data_samples' as inputs, which is image shape　actually.
        all_cls_scores, all_mask_preds = self(x, batch_data_samples)
        mask_cls_results = all_cls_scores[-1]
        mask_pred_results = all_mask_preds[-1]

        # upsample masks
        img_shape = batch_img_metas[0]['batch_input_shape']
        mask_pred_results = F.interpolate(
            mask_pred_results,
            size=img_shape,
            mode='bilinear',
            align_corners=False)

        # semantic inference
        cls_score = F.softmax(mask_cls_results, dim=-1)[..., :-1]
        mask_pred = mask_pred_results.sigmoid()
        seg_logits = torch.einsum('bqc,bqhw->bchw', cls_score, mask_pred)
        return seg_logits
//  [python]
def forward(self, input):
        # names = list(input.keys())
        input = list(input.values())

        output1 = self.output1(input[0])
        output2 = self.output2(input[1])
        output3 = self.output3(input[2])

        up3 = F.upsample(
            output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.upsample(
            output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        out = [output1, output2, output3]
        return out
//  [python]
def forward(self, input):
        # names = list(input.keys())
        input = list(input.values())

        output1 = self.output1(input[0])
        output2 = self.output2(input[1])
        output3 = self.output3(input[2])

        up3 = F.interpolate(
            output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.interpolate(
            output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        out = [output1, output2, output3]
        return out
//  [python]
def viz_weight(self, weight, img_bg=None, ii=0, jj=0, write=True):
        weight = weight.permute(1,2,0)
        # Calculate the max norm of the weight image
        #weight_norm_max = weight.norm(dim=-1).max().item()
        #ic(weight_norm_max)
        #weight = weight / weight_norm_max
        weight_x = weight[...,0][None,None]
        weight_y = weight[...,1][None,None]
        # TODO should we multiply by 8 as in upflow8?
        new_size =(self.img_height, self.img_width)
        viz_weight_x = t.nn.functional.interpolate(weight_x, size=new_size, mode="bilinear", align_corners=True).squeeze()
        viz_weight_y = t.nn.functional.interpolate(weight_y, size=new_size, mode="bilinear", align_corners=True).squeeze()
        viz_weight_x = (viz_weight_x.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit
        viz_weight_y = (viz_weight_y.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit
        viz_weight_x = cv2.applyColorMap(viz_weight_x, cv2.COLORMAP_MAGMA)
        viz_weight_y = cv2.applyColorMap(viz_weight_y, cv2.COLORMAP_MAGMA)
        if img_bg is not None:
            img_bg = img_bg.permute(1,2,0).cpu().numpy()
            viz_weight_x = cv2.addWeighted(img_bg, 0.5, viz_weight_x, 0.5, 0)
            viz_weight_y = cv2.addWeighted(img_bg, 0.5, viz_weight_y, 0.5, 0)
        name_x = f'Weights X from {ii} to {jj}'
        name_y = f'Weights Y from {ii} to {jj}'
        cv2.imshow(name_x, viz_weight_x)
        cv2.imshow(name_y, viz_weight_y)
        if write:
            cv2.imwrite(name_x+".png", viz_weight_x)
            cv2.imwrite(name_y+".png", viz_weight_y)
//  [python]
def viz_weight(self, weight, img_bg=None, ii=0, jj=0, write=True):
        weight = weight.permute(1,2,0)
        # Calculate the max norm of the weight image
        #weight_norm_max = weight.norm(dim=-1).max().item()
        #ic(weight_norm_max)
        #weight = weight / weight_norm_max
        weight_x = weight[...,0][None,None]
        weight_y = weight[...,1][None,None]
        # TODO should we multiply by 8 as in upflow8?
        new_size =(self.img_height, self.img_width)
        viz_weight_x = F.interpolate(weight_x, size=new_size, mode="bilinear", align_corners=True).squeeze()
        viz_weight_y = F.interpolate(weight_y, size=new_size, mode="bilinear", align_corners=True).squeeze()
        viz_weight_x = (viz_weight_x.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit
        viz_weight_y = (viz_weight_y.to(torch.float).cpu().numpy()*255).astype(np.uint8) # scaling from 16bit to 8bit
        viz_weight_x = cv2.applyColorMap(viz_weight_x, cv2.COLORMAP_MAGMA)
        viz_weight_y = cv2.applyColorMap(viz_weight_y, cv2.COLORMAP_MAGMA)
        if img_bg is not None:
            img_bg = img_bg.permute(1,2,0).cpu().numpy()
            viz_weight_x = cv2.addWeighted(img_bg, 0.5, viz_weight_x, 0.5, 0)
            viz_weight_y = cv2.addWeighted(img_bg, 0.5, viz_weight_y, 0.5, 0)
        name_x = f'Weights X from {ii} to {jj}'
        name_y = f'Weights Y from {ii} to {jj}'
        cv2.imshow(name_x, viz_weight_x)
        cv2.imshow(name_y, viz_weight_y)
        if write:
            cv2.imwrite(name_x+".png", viz_weight_x)
            cv2.imwrite(name_y+".png", viz_weight_y)
//  [python]
def resample_abs_pos_embed_nhwc(
        posemb,
        new_size: List[int],
        interpolation: str = 'bicubic',
        antialias: bool = True,
        verbose: bool = False,
):
    if new_size[0] == posemb.shape[-3] and new_size[1] == posemb.shape[-2]:
        return posemb

    orig_dtype = posemb.dtype
    posemb = posemb.float()
    # do the interpolation
    posemb = posemb.reshape(1, posemb.shape[-3], posemb.shape[-2], posemb.shape[-1]).permute(0, 3, 1, 2)
    posemb = t.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
    posemb = posemb.permute(0, 2, 3, 1).to(orig_dtype)

    if not torch.jit.is_scripting() and verbose:
        _logger.info(f'Resized position embedding: {posemb.shape[-3:-1]} to {new_size}.')

    return posemb
//  [python]
def resample_abs_pos_embed_nhwc(
        posemb,
        new_size: List[int],
        interpolation: str = 'bicubic',
        antialias: bool = True,
        verbose: bool = False,
):
    if new_size[0] == posemb.shape[-3] and new_size[1] == posemb.shape[-2]:
        return posemb

    orig_dtype = posemb.dtype
    posemb = posemb.float()
    # do the interpolation
    posemb = posemb.reshape(1, posemb.shape[-3], posemb.shape[-2], posemb.shape[-1]).permute(0, 3, 1, 2)
    posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
    posemb = posemb.permute(0, 2, 3, 1).to(orig_dtype)

    if not torch.jit.is_scripting() and verbose:
        _logger.info(f'Resized position embedding: {posemb.shape[-3:-1]} to {new_size}.')

    return posemb
//  [python]
def resample_abs_pos_embed(
        posemb,
        new_size: List[int],
        old_size: Optional[List[int]] = None,
        num_prefix_tokens: int = 1,
        interpolation: str = 'bicubic',
        antialias: bool = True,
        verbose: bool = False,
):
    # sort out sizes, assume square if old size not provided
    num_pos_tokens = posemb.shape[1]
    num_new_tokens = new_size[0] * new_size[1] + num_prefix_tokens
    if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]:
        return posemb

    if old_size is None:
        hw = int(math.sqrt(num_pos_tokens - num_prefix_tokens))
        old_size = hw, hw

    if num_prefix_tokens:
        posemb_prefix, posemb = posemb[:, :num_prefix_tokens], posemb[:, num_prefix_tokens:]
    else:
        posemb_prefix, posemb = None, posemb

    # do the interpolation
    embed_dim = posemb.shape[-1]
    orig_dtype = posemb.dtype
    posemb = posemb.float()  # interpolate needs float32
    posemb = posemb.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2)
    posemb = t.interpolate(posemb, size=new_size, mode=interpolation, align_corners=not antialias)
    posemb = posemb.permute(0, 2, 3, 1).reshape(1, -1, embed_dim)
    posemb = posemb.to(orig_dtype)

    # add back extra (class, etc) prefix tokens
    if posemb_prefix is not None:
        posemb = torch.cat([posemb_prefix, posemb], dim=1)

    if not torch.jit.is_scripting() and verbose:
        _logger.info(f'Resized position embedding: {old_size} to {new_size}.')

    return posemb
//  [python]
def resample_abs_pos_embed(
        posemb,
        new_size: List[int],
        old_size: Optional[List[int]] = None,
        num_prefix_tokens: int = 1,
        interpolation: str = 'bicubic',
        antialias: bool = True,
        verbose: bool = False,
):
    # sort out sizes, assume square if old size not provided
    num_pos_tokens = posemb.shape[1]
    num_new_tokens = new_size[0] * new_size[1] + num_prefix_tokens
    if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]:
        return posemb

    if old_size is None:
        hw = int(math.sqrt(num_pos_tokens - num_prefix_tokens))
        old_size = hw, hw

    if num_prefix_tokens:
        posemb_prefix, posemb = posemb[:, :num_prefix_tokens], posemb[:, num_prefix_tokens:]
    else:
        posemb_prefix, posemb = None, posemb

    # do the interpolation
    embed_dim = posemb.shape[-1]
    orig_dtype = posemb.dtype
    posemb = posemb.float()  # interpolate needs float32
    posemb = posemb.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2)
    posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
    posemb = posemb.permute(0, 2, 3, 1).reshape(1, -1, embed_dim)
    posemb = posemb.to(orig_dtype)

    # add back extra (class, etc) prefix tokens
    if posemb_prefix is not None:
        posemb = torch.cat([posemb_prefix, posemb], dim=1)

    if not torch.jit.is_scripting() and verbose:
        _logger.info(f'Resized position embedding: {old_size} to {new_size}.')

    return posemb
//  [python]
def resized_pos_embed(self, in_res, tgt_res, mode='bicubic'):
        """Resize the position embedding."""
        # assert L == (input_resolution // self.patch_size) ** 2 + 1
        L, D = self.positional_embedding.shape

        in_side = in_res // self.patch_size
        # tgt_side = tgt_res // self.patch_size
        tgt_side = int((tgt_res - 1)**0.5)

        cls_pos = self.positional_embedding[0].unsqueeze(0)  # 1 D
        pos_embed = self.positional_embedding[1:].reshape(
            1, in_side, in_side, D).permute(0, 3, 1, 2)  # L-1 D -> 1 D S S
        resized_pos_embed = t.interpolate(
            pos_embed,
            size=(tgt_side, tgt_side),
            mode=mode,
            align_corners=False,
        )  # 1 D S S -> 1 D S' S'
        resized_pos_embed = resized_pos_embed.squeeze(0).reshape(
            D, -1).T  # L'-1 D

        return torch.cat((cls_pos, resized_pos_embed), dim=0)
//  [python]
def resized_pos_embed(self, in_res, tgt_res, mode='bicubic'):
        """Resize the position embedding."""
        # assert L == (input_resolution // self.patch_size) ** 2 + 1
        L, D = self.positional_embedding.shape

        in_side = in_res // self.patch_size
        # tgt_side = tgt_res // self.patch_size
        tgt_side = int((tgt_res - 1)**0.5)

        cls_pos = self.positional_embedding[0].unsqueeze(0)  # 1 D
        pos_embed = self.positional_embedding[1:].reshape(
            1, in_side, in_side, D).permute(0, 3, 1, 2)  # L-1 D -> 1 D S S
        resized_pos_embed = F.interpolate(
            pos_embed,
            size=(tgt_side, tgt_side),
            mode=mode,
            align_corners=False,
        )  # 1 D S S -> 1 D S' S'
        resized_pos_embed = resized_pos_embed.squeeze(0).reshape(
            D, -1).T  # L'-1 D

        return torch.cat((cls_pos, resized_pos_embed), dim=0)
//  [python]
def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):        
    # interpolate position embedding
    embedding_size = pos_embed_checkpoint.shape[-1]
    num_patches = visual_encoder.patch_embed.num_patches
    num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches
    # height (== width) for the checkpoint position embedding
    orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
    # height (== width) for the new position embedding
    new_size = int(num_patches ** 0.5)

    if orig_size!=new_size:
        # class_token and dist_token are kept unchanged
        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
        # only the position tokens are interpolated
        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
        pos_tokens = t.nn.functional.interpolate(
            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
        print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))
        
        return new_pos_embed    
    else:
        return pos_embed_checkpoint
//  [python]
def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):        
    # interpolate position embedding
    embedding_size = pos_embed_checkpoint.shape[-1]
    num_patches = visual_encoder.patch_embed.num_patches
    num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches
    # height (== width) for the checkpoint position embedding
    orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
    # height (== width) for the new position embedding
    new_size = int(num_patches ** 0.5)

    if orig_size!=new_size:
        # class_token and dist_token are kept unchanged
        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
        # only the position tokens are interpolated
        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
        pos_tokens = torch.nn.functional.interpolate(
            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
        print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))
        
        return new_pos_embed    
    else:
        return pos_embed_checkpoint
//  [python]
def forward(self, loadsize, image, semantic, flow, conf, edge, backmask_in):
        # 1 x 69 x 256 x 512
        
        input = torch.cat([image, semantic, flow, conf, edge, backmask_in], dim=1)

        down1 = self.model_down_input(input)

        down2 = self.downconv1(down1)

        down3 = self.downconv2(down2)


        down4 = self.downconv3(down3)

        feat = self.model_res_up(self.model_res_down(down4))

        up4 = F.upsample(feat, scale_factor=2, mode='bilinear')
        
        up4_concat = torch.cat([up4, down3], dim=1)
 
        up4_conv = self.upconv1(up4_concat)


        flow3 = self.model_final_flow_3(up4_conv)

        flow3_up = F.upsample(flow3, scale_factor=2, mode='bilinear')

        up3 = F.upsample(up4_conv, scale_factor=2, mode='bilinear')

        up3_concat = torch.cat([up3, down2, flow3_up], dim=1)

        up3_conv = self.upconv2(up3_concat)

        flow2 = self.model_final_flow_2(up3_conv)

        flow2_up = F.upsample(flow2, scale_factor=2, mode='bilinear')
        
        up2 = F.upsample(up3_conv, scale_factor=2, mode='bilinear')
       
        up2_concat = torch.cat([up2, down1, flow2_up], dim=1)
        up2_conv = self.upconv3(up2_concat)

        flow1 = self.model_final_flow_1(up2_conv)

        return [flow1, flow2, flow3], up2_conv, input
//  [python]
def forward(self, loadsize, image, semantic, flow, conf, edge, backmask_in):
        # 1 x 69 x 256 x 512
        
        input = torch.cat([image, semantic, flow, conf, edge, backmask_in], dim=1)

        down1 = self.model_down_input(input)

        down2 = self.downconv1(down1)

        down3 = self.downconv2(down2)


        down4 = self.downconv3(down3)

        feat = self.model_res_up(self.model_res_down(down4))

        up4 = F.interpolate(feat, scale_factor=2, mode='bilinear', align_corners=False)
        
        up4_concat = torch.cat([up4, down3], dim=1)
 
        up4_conv = self.upconv1(up4_concat)


        flow3 = self.model_final_flow_3(up4_conv)

        flow3_up = F.interpolate(flow3, scale_factor=2, mode='bilinear', align_corners=False)

        up3 = F.interpolate(up4_conv, scale_factor=2, mode='bilinear', align_corners=False)

        up3_concat = torch.cat([up3, down2, flow3_up], dim=1)

        up3_conv = self.upconv2(up3_concat)

        flow2 = self.model_final_flow_2(up3_conv)

        flow2_up = F.interpolate(flow2, scale_factor=2, mode='bilinear', align_corners=False)
        
        up2 = F.interpolate(up3_conv, scale_factor=2, mode='bilinear', align_corners=False)
       
        up2_concat = torch.cat([up2, down1, flow2_up], dim=1)
        up2_conv = self.upconv3(up2_concat)

        flow1 = self.model_final_flow_1(up2_conv)

        return [flow1, flow2, flow3], up2_conv, input
//  [python]
def encode(self, vae, pixels, mask, grow_mask_by=6):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        mask = t.nn.functional.upsample(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

        pixels = pixels.clone()
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]

        #grow mask by a few pixels to keep things seamless in latent space
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)

        m = (1.0 - mask.round()).squeeze(1)
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
            pixels[:,:,:,i] *= m
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )
//  [python]
def encode(self, vae, pixels, mask, grow_mask_by=6):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

        pixels = pixels.clone()
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]

        #grow mask by a few pixels to keep things seamless in latent space
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)

        m = (1.0 - mask.round()).squeeze(1)
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
            pixels[:,:,:,i] *= m
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )