diff --git a/SDMatte/__init__.py b/SDMatte/__init__.py new file mode 100644 index 0000000..006f212 --- /dev/null +++ b/SDMatte/__init__.py @@ -0,0 +1,2 @@ +# 使 SDMatte 目录成为可导入包 + diff --git a/SDMatte/modeling/SDMatte/__init__.py b/SDMatte/modeling/SDMatte/__init__.py new file mode 100644 index 0000000..a67c8be --- /dev/null +++ b/SDMatte/modeling/SDMatte/__init__.py @@ -0,0 +1 @@ +from .meta_arch import * diff --git a/SDMatte/modeling/SDMatte/meta_arch.py b/SDMatte/modeling/SDMatte/meta_arch.py new file mode 100644 index 0000000..71c69be --- /dev/null +++ b/SDMatte/modeling/SDMatte/meta_arch.py @@ -0,0 +1,235 @@ +import torch +from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig +import torch.nn as nn +import torch.nn.functional as F +from diffusers import DDIMScheduler, AutoencoderKL, UNet2DConditionModel +from diffusers.models.embeddings import get_timestep_embedding +from ...utils import replace_unet_conv_in, replace_attention_mask_method, add_aux_conv_in +from ...utils.replace import CustomUNet +import random +import os + +# 解决离线本地目录层级差异，如存在 "subdir/subdir/config.json" 的情况 +def _resolve_nested_dir(base_dir: str, subdir: str, config_filename: str) -> str: + direct = os.path.join(base_dir, subdir) + nested = os.path.join(base_dir, subdir, subdir) + if os.path.exists(os.path.join(direct, config_filename)): + return direct + if os.path.exists(os.path.join(nested, config_filename)): + return nested + return direct + +AUX_INPUT_DIT = { + "auto_mask": "auto_coords", + "point_mask": "point_coords", + "bbox_mask": "bbox_coords", + "mask": "mask_coords", + "trimap": "trimap_coords", +} + +class SDMatte(nn.Module): + def __init__( + self, + pretrained_model_name_or_path, + conv_scale=3, + num_inference_steps=1, + aux_input="bbox_mask", + use_aux_input=False, + use_coor_input=True, + use_dis_loss=True, + use_attention_mask=True, + use_encoder_attention_mask=False, + add_noise=False, + attn_mask_aux_input=["point_mask", "bbox_mask", "mask"], + aux_input_list=["point_mask", "bbox_mask", "mask"], + use_encoder_hidden_states=True, + residual_connection=False, + use_attention_mask_list=[True, True, True], + use_encoder_hidden_states_list=[True, True, True], + load_weight = True, + ): + super().__init__() + self.init_submodule(pretrained_model_name_or_path, load_weight) + self.num_inference_steps = num_inference_steps + self.aux_input = aux_input + self.use_aux_input = use_aux_input + self.use_coor_input = use_coor_input + self.use_dis_loss = use_dis_loss + self.use_attention_mask = use_attention_mask + self.use_encoder_attention_mask = use_encoder_attention_mask + self.add_noise = add_noise + self.attn_mask_aux_input = attn_mask_aux_input + self.aux_input_list = aux_input_list + self.use_encoder_hidden_states = use_encoder_hidden_states + if use_encoder_hidden_states: + self.unet = add_aux_conv_in(self.unet) + if not add_noise: + conv_scale -= 1 + if not use_aux_input: + conv_scale -= 1 + if conv_scale > 1: + self.unet = replace_unet_conv_in(self.unet, conv_scale) + replace_attention_mask_method(self.unet, residual_connection) + self.text_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + self.unet.train() + self.unet.use_attention_mask_list = use_attention_mask_list + self.unet.use_encoder_hidden_states_list = use_encoder_hidden_states_list + + def init_submodule(self, pretrained_model_name_or_path, load_weight): + if load_weight: + text_dir = _resolve_nested_dir(pretrained_model_name_or_path, "text_encoder", "config.json") + vae_dir = _resolve_nested_dir(pretrained_model_name_or_path, "vae", "config.json") + unet_dir = _resolve_nested_dir(pretrained_model_name_or_path, "unet", "config.json") + sched_dir = _resolve_nested_dir(pretrained_model_name_or_path, "scheduler", "scheduler_config.json") + tok_dir = _resolve_nested_dir(pretrained_model_name_or_path, "tokenizer", "tokenizer_config.json") + + self.text_encoder = CLIPTextModel.from_pretrained(text_dir) + self.vae = AutoencoderKL.from_pretrained(vae_dir) + self.unet = CustomUNet.from_pretrained( + unet_dir, low_cpu_mem_usage=True, ignore_mismatched_sizes=False + ) + self.noise_scheduler = DDIMScheduler.from_pretrained(sched_dir) + self.tokenizer = CLIPTokenizer.from_pretrained(tok_dir) + else: + text_dir = _resolve_nested_dir(pretrained_model_name_or_path, "text_encoder", "config.json") + text_config = CLIPTextConfig.from_pretrained(text_dir) + self.text_encoder = CLIPTextModel(text_config) + + vae_path = _resolve_nested_dir(pretrained_model_name_or_path, "vae", "config.json") + self.vae = AutoencoderKL.from_config(AutoencoderKL.load_config(vae_path)) + + unet_path = _resolve_nested_dir(pretrained_model_name_or_path, "unet", "config.json") + self.unet = CustomUNet.from_config( + CustomUNet.load_config(unet_path), + low_cpu_mem_usage=True, + ignore_mismatched_sizes=False + ) + + scheduler_path = os.path.join(_resolve_nested_dir(pretrained_model_name_or_path, "scheduler", "scheduler_config.json"), "scheduler_config.json") + self.noise_scheduler = DDIMScheduler.from_config(DDIMScheduler.load_config(scheduler_path)) + + tok_dir = _resolve_nested_dir(pretrained_model_name_or_path, "tokenizer", "tokenizer_config.json") + self.tokenizer = CLIPTokenizer.from_pretrained(tok_dir) + + + def forward(self, data): + rgb = data["image"].cuda() + B = rgb.shape[0] + + if self.aux_input is None and self.training: + aux_input_type = random.choice(self.aux_input_list) + elif self.aux_input is None: + aux_input_type = "point_mask" + else: + aux_input_type = self.aux_input + + # get aux input latent + if self.use_aux_input: + aux_input = data[aux_input_type].cuda() + aux_input = aux_input.repeat(1, 3, 1, 1) + aux_input_h = self.vae.encoder(aux_input.to(rgb.dtype)) + aux_input_moments = self.vae.quant_conv(aux_input_h) + aux_input_mean, _ = torch.chunk(aux_input_moments, 2, dim=1) + aux_input_latent = aux_input_mean * self.vae.config.scaling_factor + else: + aux_input_latent = None + + # get aux coordinate + coor_name = AUX_INPUT_DIT[aux_input_type] + coor = data[coor_name].cuda() + if coor_name == "point_coords": + N = coor.shape[1] + for i in range(N, 1680): + if 1680 % i == 0: + num_channels = 1680 // i + pad_size = i - N + padding = torch.zeros((B, pad_size), dtype=coor.dtype, device=coor.device) + coor = torch.cat([coor, padding], dim=1) + zero_coor = torch.zeros((B, pad_size + N), dtype=coor.dtype, device=coor.device) + break + if self.use_coor_input: + coor = get_timestep_embedding( + coor.flatten(), + num_channels, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + else: + coor = get_timestep_embedding( + zero_coor.flatten(), + num_channels, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + added_cond_kwargs = {"point_coords": coor} + else: + if self.use_coor_input: + added_cond_kwargs = {"bbox_mask_coords": coor} + else: + coor = torch.tensor([[0, 0, 1, 1]] * B).cuda() + added_cond_kwargs = {"bbox_mask_coords": coor} + + # get attention mask + if self.use_attention_mask and aux_input_type in self.attn_mask_aux_input: + attention_mask = data[aux_input_type].cuda() + attention_mask = (attention_mask + 1) / 2 + attention_mask = F.interpolate(attention_mask, scale_factor=1 / 8, mode="nearest") + attention_mask = attention_mask.flatten(start_dim=1) + else: + attention_mask = None + + # encode rgb to latents + rgb_h = self.vae.encoder(rgb) + rgb_moments = self.vae.quant_conv(rgb_h) + rgb_mean, _ = torch.chunk(rgb_moments, 2, dim=1) + rgb_latent = rgb_mean * self.vae.config.scaling_factor + + # get encoder_hidden_states + if self.use_encoder_hidden_states and aux_input_latent is not None: + encoder_hidden_states = self.unet.aux_conv_in(aux_input_latent) + encoder_hidden_states = encoder_hidden_states.view(B, 1024, -1) + encoder_hidden_states = encoder_hidden_states.permute(0, 2, 1) + + if "caption" in data: + prompt = data["caption"] + else: + prompt = [""] * B + prompt = [prompt] if isinstance(prompt, str) else prompt + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to("cuda") + text_embed = self.text_encoder(text_input_ids)[0] + encoder_hidden_states_2 = text_embed + + # get class_label + is_trans = data["is_trans"].cuda() + trans = 1 - is_trans + + # get timesteps + timestep = torch.tensor([1], device="cuda").long() + + # unet + unet_input = torch.cat([rgb_latent, aux_input_latent], dim=1) + label_latent = self.unet( + sample=unet_input, + trans=trans, + timestep=None, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_2=encoder_hidden_states_2, + added_cond_kwargs=added_cond_kwargs, + attention_mask=attention_mask, + ).sample + label_latent = label_latent / self.vae.config.scaling_factor + z = self.vae.post_quant_conv(label_latent) + stacked = self.vae.decoder(z) + # mean of output channels + label_mean = stacked.mean(dim=1, keepdim=True) + output = torch.clip(label_mean, -1.0, 1.0) + output = (output + 1.0) / 2.0 + return output diff --git a/SDMatte/modeling/__init__.py b/SDMatte/modeling/__init__.py new file mode 100644 index 0000000..153fb8d --- /dev/null +++ b/SDMatte/modeling/__init__.py @@ -0,0 +1,2 @@ +from .SDMatte import * +# from .LiteSDMatte import * diff --git a/SDMatte/utils/__init__.py b/SDMatte/utils/__init__.py new file mode 100644 index 0000000..2b54a1f --- /dev/null +++ b/SDMatte/utils/__init__.py @@ -0,0 +1,2 @@ +from .utils import * +from .replace import * diff --git a/SDMatte/utils/replace.py b/SDMatte/utils/replace.py new file mode 100644 index 0000000..a32347e --- /dev/null +++ b/SDMatte/utils/replace.py @@ -0,0 +1,550 @@ +import math +from typing import Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +import math +from diffusers import UNet2DConditionModel +from diffusers.models.embeddings import Timesteps, TimestepEmbedding +from diffusers.models.unets.unet_2d_blocks import ( + get_down_block, + get_up_block, + get_mid_block, +) +from diffusers.models.activations import get_activation +from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput +from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers + + +def custom_prepare_attention_mask( + self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 +) -> torch.Tensor: + r""" + Prepare the attention mask for the attention computation. + + Args: + attention_mask (`torch.Tensor`): + The attention mask to prepare. + target_length (`int`): + The target length of the attention mask. This is the length of the attention mask after padding. + batch_size (`int`): + The batch size, which is used to repeat the attention mask. + out_dim (`int`, *optional*, defaults to `3`): + The output dimension of the attention mask. Can be either `3` or `4`. + + Returns: + `torch.Tensor`: The prepared attention mask. + """ + head_size = self.heads + if attention_mask is None: + return attention_mask + + current_length: int = attention_mask.shape[-1] + if current_length != target_length: + if attention_mask.device.type == "mps": + # HACK: MPS: Does not support padding by greater than dimension of input tensor. + # Instead, we can manually construct the padding tensor. + padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length) + padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) + attention_mask = torch.cat([attention_mask, padding], dim=2) + else: + # TODO: for pipelines such as stable-diffusion, padding cross-attn mask: + # we want to instead pad by (0, remaining_length), where remaining_length is: + # remaining_length: int = target_length - current_length + # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding + B = attention_mask.shape[0] + current_size = int(math.sqrt(current_length)) + target_size = int(math.sqrt(target_length)) + assert current_size**2 == current_length, f"current_length ({current_length}) cannot be squared to an integer size" + assert target_size**2 == target_length, f"target_length ({target_length}) cannot be squared to an integer size" + attention_mask = attention_mask.view(B, -1, current_size, current_size) + attention_mask = F.interpolate(attention_mask, size=(target_size, target_size), mode="nearest") + attention_mask = attention_mask.view(B, 1, target_length) + + if out_dim == 3: + if attention_mask.shape[0] < batch_size * head_size: + attention_mask = attention_mask.repeat_interleave(head_size, dim=0) + elif out_dim == 4: + attention_mask = attention_mask.unsqueeze(1) + attention_mask = attention_mask.repeat_interleave(head_size, dim=1) + + return attention_mask + + +def custom_get_attention_scores(self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor: + r""" + Compute the attention scores. + + Args: + query (`torch.Tensor`): The query tensor. + key (`torch.Tensor`): The key tensor. + attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied. + + Returns: + `torch.Tensor`: The attention probabilities/scores. + """ + dtype = query.dtype + if self.upcast_attention: + query = query.float() + key = key.float() + + # if attention_mask is not None and len(torch.unique(attention_mask)) <= 2: + if attention_mask is not None: + baddbmm_input = attention_mask + beta = 1 + else: + baddbmm_input = torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device) + beta = 0 + + attention_scores = torch.baddbmm( + baddbmm_input, + query, + key.transpose(-1, -2), + beta=beta, + alpha=self.scale, + ) + + # if attention_mask is not None and len(torch.unique(attention_mask)) > 2: + # m = 1 - (attention_mask / -10000.0) + # attention_scores = m * attention_scores + + del baddbmm_input + + if self.upcast_softmax: + attention_scores = attention_scores.float() + + attention_probs = attention_scores.softmax(dim=-1) + del attention_scores + + attention_probs = attention_probs.to(dtype) + + return attention_probs + + +class CustomUNet(UNet2DConditionModel): + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + dropout: float = 0.0, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, + reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_dim: Optional[int] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + bbox_time_embed_dim: Optional[int] = None, + point_embeddings_input_dim: Optional[int] = None, + bbox_embeddings_input_dim: Optional[int] = None, + attention_type: str = "default", + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + use_attention_mask_list=[True, True, True], + use_encoder_hidden_states_list=[True, True, True], + ): + super().__init__() + self.use_attention_mask_list = use_attention_mask_list + self.use_encoder_hidden_states_list = use_encoder_hidden_states_list + self.sample_size = sample_size + num_attention_heads = num_attention_heads or attention_head_dim + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding) + + # time + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + self.point_embedding = TimestepEmbedding(point_embeddings_input_dim, time_embed_dim) + self.bbox_time_proj = Timesteps(bbox_time_embed_dim, flip_sin_to_cos, freq_shift) + self.bbox_embedding = TimestepEmbedding(bbox_embeddings_input_dim, time_embed_dim) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = get_mid_block( + mid_block_type, + temb_channels=blocks_time_embed_dim, + in_channels=block_out_channels[-1], + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + output_scale_factor=mid_block_scale_factor, + transformer_layers_per_block=transformer_layers_per_block[-1], + num_attention_heads=num_attention_heads[-1], + cross_attention_dim=cross_attention_dim[-1], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + mid_block_only_cross_attention=mid_block_only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[-1], + dropout=dropout, + ) + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = ( + list(reversed(transformer_layers_per_block)) + if reverse_transformer_layers_per_block is None + else reverse_transformer_layers_per_block + ) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resolution_idx=i, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps) + + self.conv_act = get_activation(act_fn) + + else: + self.conv_norm_out = None + self.conv_act = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding) + + # distillation + self.feature_map = [] + + def _get_value(self, use_list, true_value, false_value): + down_value = mid_value = up_value = false_value + + if use_list[0]: + down_value = true_value + if use_list[1]: + mid_value = true_value + if use_list[2]: + up_value = true_value + + return down_value, mid_value, up_value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + trans: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_2: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + ) -> Union[UNet2DConditionOutput, Tuple]: + default_overall_up_factor = 2**self.num_upsamplers + forward_upsample_size = False + upsample_size = None + + for dim in sample.shape[-2:]: + if dim % default_overall_up_factor != 0: + forward_upsample_size = True + break + + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + if encoder_attention_mask is not None: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + down_attn_mask, mid_attn_mask, up_attn_mask = self._get_value(self.use_attention_mask_list, attention_mask, None) + down_encoder_hidden_states, mid_encoder_hidden_states, up_encoder_hidden_states = self._get_value( + self.use_encoder_hidden_states_list, encoder_hidden_states, encoder_hidden_states_2 + ) + + # 1. time + t_emb, op_emb, aug_emb = None, None, None + + if timestep is not None: + timesteps = timestep + timesteps = timesteps.expand(sample.shape[0]) + t_emb = self.time_proj(timesteps) + t_emb = t_emb.to(dtype=sample.dtype) + + t_emb = self.time_embedding(t_emb, timestep_cond) + + # opacity + if trans is not None: + trans = trans.expand(sample.shape[0]) + op_emb = self.time_proj(trans) + op_emb = op_emb.to(dtype=sample.dtype) + + op_emb = self.time_embedding(op_emb, timestep_cond) + + if t_emb is not None and op_emb is not None: + emb = t_emb + op_emb + elif op_emb is not None: + emb = op_emb + elif t_emb is not None: + emb = t_emb + else: + raise ValueError("Missing required field: 'timestep' and 'trans'. Please ensure it is included in your input.") + + if "point_coords" in added_cond_kwargs: + coords_embeds = added_cond_kwargs.get("point_coords") + coords_embeds = coords_embeds.reshape((sample.shape[0], -1)) + coords_embeds = coords_embeds.to(emb.dtype) + aug_emb = self.point_embedding(coords_embeds) + elif "bbox_mask_coords" in added_cond_kwargs: + coords = added_cond_kwargs.get("bbox_mask_coords") + coords_embeds = self.bbox_time_proj(coords.flatten()) + coords_embeds = coords_embeds.reshape((sample.shape[0], -1)) + coords_embeds = coords_embeds.to(emb.dtype) + aug_emb = self.bbox_embedding(coords_embeds) + else: + raise ValueError(f"{self.__class__} cannot find point_coords or bbox_coords in added_cond_kwargs.") + + emb = emb + aug_emb if aug_emb is not None else emb + + # 2. pre-process + sample = self.conv_in(sample) + + # distillation + self.feature_map = [] + + # 3. down + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + if USE_PEFT_BACKEND: + scale_lora_layers(self, lora_scale) + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + additional_residuals = {} + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=down_encoder_hidden_states, + attention_mask=down_attn_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale) + + down_block_res_samples += res_samples + + self.feature_map.append(sample) + + # 4. mid + if self.mid_block is not None: + if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=mid_encoder_hidden_states, + attention_mask=mid_attn_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = self.mid_block(sample, emb) + + self.feature_map.append(sample) + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=up_encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=up_attn_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + ) + + self.feature_map.append(sample) + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if USE_PEFT_BACKEND: + unscale_lora_layers(self, lora_scale) + + return UNet2DConditionOutput(sample=sample) diff --git a/SDMatte/utils/utils.py b/SDMatte/utils/utils.py new file mode 100644 index 0000000..dcaa0f1 --- /dev/null +++ b/SDMatte/utils/utils.py @@ -0,0 +1,83 @@ +import os +import json +import torch.nn as nn +from torch.nn import Conv2d +from torch.nn.parameter import Parameter +from diffusers.models.attention_processor import Attention, AttnProcessor +from .replace import custom_prepare_attention_mask, custom_get_attention_scores +import cv2 +import torch +import numpy as np + + +def replace_unet_conv_in(unet, num): + # replace the first layer to accept 8 in_channels + _weight = unet.conv_in.weight.clone() # [320, 4, 3, 3] + _bias = unet.conv_in.bias.clone() # [320] + _weight = _weight.repeat((1, num, 1, 1)) # Keep selected channel(s) + # half the activation magnitude + _weight = _weight / num + # new conv_in channel + _n_convin_out_channel = unet.conv_in.out_channels + _new_conv_in = Conv2d(4 * num, _n_convin_out_channel, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + _new_conv_in.weight = Parameter(_weight) + _new_conv_in.bias = Parameter(_bias) + unet.conv_in = _new_conv_in + print("Unet conv_in layer is replaced") + # replace config + unet.config["in_channels"] = 4 * num + print("Unet config is updated") + return unet + + +def add_aux_conv_in(unet): + aux_conv_in = nn.Conv2d(in_channels=4, out_channels=1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + aux_conv_in.weight.data[:320, :, :, :] = unet.conv_in.weight.data.clone() + aux_conv_in.weight.data[320:, :, :, :] = 0.0 + aux_conv_in.bias.data[:320] = unet.conv_in.bias.data.clone() + aux_conv_in.bias.data[320:] = 0.0 + unet.aux_conv_in = aux_conv_in + print("add aux_conv_in layer for unet") + return unet + + +def replace_attention_mask_method(module, residual_connection): + if isinstance(module, Attention): + module.processor = AttnProcessor() + if hasattr(module, "prepare_attention_mask"): + module.prepare_attention_mask = custom_prepare_attention_mask.__get__(module) + if hasattr(module, "cross_attention_dim") and module.cross_attention_dim == 320: + module.residual_connection = residual_connection + if hasattr(module, "get_attention_scores"): + module.get_attention_scores = custom_get_attention_scores.__get__(module) + + # 递归遍历所有子模块 + for child_name, child_module in module.named_children(): + replace_attention_mask_method(child_module, residual_connection) + + +erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1, 30)] + + +def get_unknown_tensor_from_pred(pred, rand_width=30, train_mode=True): + ### pred: N, 1 ,H, W + N, C, H, W = pred.shape + + pred = pred.data.cpu().numpy() + uncertain_area = np.ones_like(pred, dtype=np.uint8) + uncertain_area[pred < 1.0 / 255.0] = 0 + uncertain_area[pred > 1 - 1.0 / 255.0] = 0 + + for n in range(N): + uncertain_area_ = uncertain_area[n, 0, :, :] # H, W + if train_mode: + width = np.random.randint(1, rand_width) + else: + width = rand_width // 2 + uncertain_area_ = cv2.dilate(uncertain_area_, erosion_kernels[width]) + uncertain_area[n, 0, :, :] = uncertain_area_ + + weight = np.zeros_like(uncertain_area) + weight[uncertain_area == 1] = 1 + weight = torch.from_numpy(weight).float().cuda() + return weight