Add files via upload
This commit is contained in:
AI Lab
@@ -0,0 +1,2 @@
|
||||
# 使 SDMatte 目录成为可导入包
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .meta_arch import *
|
||||
@@ -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
|
||||
@@ -0,0 +1,2 @@
|
||||
from .SDMatte import *
|
||||
# from .LiteSDMatte import *
|
||||
@@ -0,0 +1,2 @@
|
||||
from .utils import *
|
||||
from .replace import *
|
||||
@@ -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)
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user