From 75e6c6c5b24b67b286dd39b6ad9b88e88ce00287 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Thu, 1 Jan 2026 19:07:44 -0800 Subject: [PATCH] Delete SDMatte directory --- SDMatte/__init__.py | 2 - SDMatte/modeling/SDMatte/__init__.py | 1 - SDMatte/modeling/SDMatte/meta_arch.py | 235 ----------- SDMatte/modeling/__init__.py | 2 - SDMatte/utils/__init__.py | 2 - SDMatte/utils/replace.py | 550 -------------------------- SDMatte/utils/utils.py | 83 ---- 7 files changed, 875 deletions(-) delete mode 100644 SDMatte/__init__.py delete mode 100644 SDMatte/modeling/SDMatte/__init__.py delete mode 100644 SDMatte/modeling/SDMatte/meta_arch.py delete mode 100644 SDMatte/modeling/__init__.py delete mode 100644 SDMatte/utils/__init__.py delete mode 100644 SDMatte/utils/replace.py delete mode 100644 SDMatte/utils/utils.py diff --git a/SDMatte/__init__.py b/SDMatte/__init__.py deleted file mode 100644 index 006f212..0000000 --- a/SDMatte/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# 使 SDMatte 目录成为可导入包 - diff --git a/SDMatte/modeling/SDMatte/__init__.py b/SDMatte/modeling/SDMatte/__init__.py deleted file mode 100644 index a67c8be..0000000 --- a/SDMatte/modeling/SDMatte/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .meta_arch import * diff --git a/SDMatte/modeling/SDMatte/meta_arch.py b/SDMatte/modeling/SDMatte/meta_arch.py deleted file mode 100644 index 71c69be..0000000 --- a/SDMatte/modeling/SDMatte/meta_arch.py +++ /dev/null @@ -1,235 +0,0 @@ -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 deleted file mode 100644 index 153fb8d..0000000 --- a/SDMatte/modeling/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .SDMatte import * -# from .LiteSDMatte import * diff --git a/SDMatte/utils/__init__.py b/SDMatte/utils/__init__.py deleted file mode 100644 index 2b54a1f..0000000 --- a/SDMatte/utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .utils import * -from .replace import * diff --git a/SDMatte/utils/replace.py b/SDMatte/utils/replace.py deleted file mode 100644 index a32347e..0000000 --- a/SDMatte/utils/replace.py +++ /dev/null @@ -1,550 +0,0 @@ -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 deleted file mode 100644 index dcaa0f1..0000000 --- a/SDMatte/utils/utils.py +++ /dev/null @@ -1,83 +0,0 @@ -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