''' Use pretrained instruct pix2pix model but add additional channels for reference modification ''' import torch from .diffusion import DDIMLDMTextTraining from einops import rearrange from modules.video_unet_temporal.resnet import InflatedConv3d from safetensors.torch import load_file import torch.nn.functional as F from torch import nn import cv2 from torch.hub import download_url_to_file class MLP(nn.Module): def __init__(self): super(MLP, self).__init__() self.fc1 = nn.Linear(3072, 4096) self.fc2 = nn.Linear(4096, 4096) self.fc3 = nn.Linear(4096, 4096) self.fc4 = nn.Linear(4096, 2304) self.leaky_relu = nn.LeakyReLU(negative_slope=0.01) # 设置Leaky ReLU的负斜率 def forward(self, x): x = self.leaky_relu(self.fc1(x)) x = self.leaky_relu(self.fc2(x)) x = self.leaky_relu(self.fc3(x)) x = self.fc4(x) return x # class CombineMLP(nn.Module): # def __init__(self, input_dim=128, output_dim=64, hidden_dim=128): # """ # 构造一个 5 层 MLP 网络。 # :param input_dim: 输入的特征维度，默认 128 # :param output_dim: 输出的特征维度，默认 64 # :param hidden_dim: 隐藏层维度，默认 128 # """ # super(CombineMLP, self).__init__() # # 定义 5 层 MLP # self.fc1 = nn.Linear(input_dim, hidden_dim) #() # self.fc2 = nn.Linear(hidden_dim, hidden_dim) # self.fc3 = nn.Linear(hidden_dim, hidden_dim) # self.fc4 = nn.Linear(hidden_dim, hidden_dim) # self.fc5 = nn.Linear(hidden_dim, output_dim) # 最后一层映射到 64 # # 定义激活函数 # # self.activation = nn.ReLU() # self.activation = nn.LeakyReLU(negative_slope=0.01) # 默认负斜率为 0.01 # def forward(self, x1, x2): # """ # 前向传播，支持两个输入 x1 和 x2 # :param x1: 第一个输入，形状 (B, 64) # :param x2: 第二个输入，形状 (B, 64) # :return: 输出特征，形状 (B, 64) # """ # # 将两个输入拼接在一起 # x = torch.cat([x1, x2], dim=-1) # 拼接后形状为 (B, 128) # # 依次通过 5 层 MLP 和激活函数 # x = self.activation(self.fc1(x)) # x = self.activation(self.fc2(x)) # x = self.activation(self.fc3(x)) # x = self.activation(self.fc4(x)) # x = self.fc5(x) # 最后一层不使用激活函数（根据需求） # return x class CombineMLP(nn.Module): def __init__(self, input_dim=4*64*64*2, output_dim=4*64*64, hidden_dim=128, num_layers=5): """ 构造一个 5 层 MLP 网络。 :param input_dim: 输入的特征维度，默认 128 :param output_dim: 输出的特征维度，默认 64 :param hidden_dim: 隐藏层维度，默认 128 """ super(CombineMLP, self).__init__() # 创建多个隐藏层 layers = [] for i in range(num_layers - 1): # 生成 num_layers-1 个隐藏层 layers.append(nn.Linear(input_dim if i == 0 else hidden_dim, hidden_dim)) layers.append(nn.ReLU()) # 输出层 layers.append(nn.Linear(hidden_dim, output_dim)) # 将层组合成一个模块 self.mlp = nn.Sequential(*layers) def forward(self, x1, x2): """ 前向传播，支持两个输入 x1 和 x2 :param x1: 第一个输入，形状 (1,16,4,64,64) :param x2: 第二个输入，形状 (1,16,4,64,64) :return: 输出特征，形状 (1,16,4,64,64) """ # import pdb; pdb.set_trace() # 将两个输入拼接在一起 x = torch.cat([x1, x2], dim=2) # 拼接后形状为 (1,16,8,64,64) x = torch.flatten(x, start_dim=2) # Flatten to shape (batch_size, 16, 8*64*64) x = self.mlp(x) # Apply MLP 1,16,16384 x = x.reshape(x.size(0), x.size(1), 4, 64, 64) # Reshape back to (1, 16, 4, 64, 64) return x class HDRCtrlModeltmp(nn.Module): def __init__(self): super(HDRCtrlModel, self).__init__() # 定义卷积层 self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=4, padding=1) self.conv_layer2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=2, padding=1) # 定义 MLP 模型 self.mlp = MLP() def decompose_hdr(self, hdr_latents): batch_size, channels, height, width = hdr_latents.shape device = hdr_latents.device # 获取设备信息 # 生成 4×4 掩码 (batch_size, 1, 4, 4) mask_small = torch.rand(batch_size, 1, 4, 4, device=device) # 从均匀分布生成随机掩码 # 将掩码调整为与输入相同的大小 (batch_size, 1, height, width) mask = torch.nn.functional.interpolate(mask_small, size=(height, width), mode='bilinear', align_corners=False) # 保持连续值，不进行二值化 #! 注意此步操作, 注意可视化 random mask的结果... 首先可以可视化mask, 其次可视化 mask = mask.expand(-1, channels, -1, -1) # 扩展掩码通道数以匹配 hdr_latents 的形状 # 应用 mask 生成 L1 和 L2 hdr_latents_1 = hdr_latents * mask # L1 = 掩码部分 hdr_latents_2 = hdr_latents * (1 - mask) # L2 = 非掩码部分 return hdr_latents_1, hdr_latents_2 def forward(self, hdr_latents): # import pdb; pdb.set_trace() # todo: mask get hdr1, hdr2; input hdr_latents(实际上暂时是ldr) # 输入的形状为 (1, 16, 3, 256, 512)，去掉多余的维度 # import pdb; pdb.set_trace() hdr_latents = hdr_latents.squeeze(0) # 变成 (16, 3, 256, 512) batch_size = hdr_latents.shape[0] # 转换为 NCHW 形式： (batch, channels, height, width) 输入之前numpy2tensor已经permute过了 # hdr_latents = hdr_latents.permute(0, 3, 1, 2) #! 注意一下to tensor? (如何进行归一化的) 的时候已经 # 进行卷积操作 conv_output = self.conv_layer1(hdr_latents) #! 注意更改此处卷积!!! conv_output = self.conv_layer2(conv_output) # (16, 3, 32, 64) # 截取前 32 列，得到最终形状 (16, 3, 32, 32) hdr_latents = conv_output[:, :, :, :32] # todo: decompose hdr hdr_latents_1, hdr_latents_2 = self.decompose_hdr(hdr_latents) # [16, 3, 32, 32], [16, 3, 32, 32] # 将输出展平，准备输入到 MLP 中 hdr_latents = hdr_latents.reshape(hdr_latents.size(0), -1) # [16, 3072] hdr_latents_1 = hdr_latents_1.reshape(hdr_latents_1.size(0), -1) # [16, 3072] hdr_latents_2 = hdr_latents_2.reshape(hdr_latents_2.size(0), -1) # 传递给 MLP hdr_latents = self.mlp(hdr_latents) #(16, 2304) 3072 -> 2304 hdr_latents_1 = self.mlp(hdr_latents_1) hdr_latents_2 = self.mlp(hdr_latents_2) # 重新调整输出的形状 hdr_latents = hdr_latents.reshape(batch_size, 3, 768) # reshape 输出为 (16, 3, 768) hdr_latents_1 = hdr_latents_1.reshape(batch_size, 3, 768) hdr_latents_2 = hdr_latents_2.reshape(batch_size, 3, 768) return hdr_latents, hdr_latents_1, hdr_latents_2 class HDRCtrlModel(nn.Module): def __init__(self): super(HDRCtrlModel, self).__init__() # 定义卷积层 self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=4, padding=1) self.conv_layer2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=2, padding=1) # 定义 MLP 模型 self.mlp = MLP() def decompose_hdr(self, hdr_latents): # hdr_latents: 16,3,32,32 可以可视化一下这部分代码... batch_size, channels, height, width = hdr_latents.shape device = hdr_latents.device # 获取设备信息 # 生成 4×4 掩码 (batch_size, 1, 4, 4) mask_small = torch.rand(batch_size, 1, 4, 4, device=device) # 从均匀分布生成随机掩码 threshold = 0.5 # 调节阈值，增加黑色部分比例 mask_small = (mask_small > threshold).float() # 将掩码调整为与输入相同的大小 (batch_size, 1, height, width) 16,1,32,32 mask = torch.nn.functional.interpolate(mask_small, size=(height, width), mode='bilinear', align_corners=False) # import pdb; pdb.set_trace() # 保持连续值，不进行二值化 #! 注意此步操作, 注意可视化 random mask的结果... 首先可以可视化mask, 其次可视化 mask = mask.expand(-1, channels, -1, -1) # 扩展掩码通道数以匹配 hdr_latents 的形状 # 应用 mask 生成 L1 和 L2 hdr_latents_1 = hdr_latents * mask # L1 = 掩码部分 hdr_latents_2 = hdr_latents * (1 - mask) # L2 = 非掩码部分 return hdr_latents_1, hdr_latents_2 def blur_image(self, hdr_latents): # 高斯模糊, 输入 (16,3,256,256) processed_images = [] kernel_size = (15, 15) sigmaX = 10 # 对每张图像进行处理 for i in range(hdr_latents.size(0)): # 遍历16张图像 # 获取第i张图像 image = hdr_latents[i].permute(1, 2, 0).cpu().numpy() # 将形状变为 (256, 256, 3) # 进行高斯模糊 blurred_image = cv2.GaussianBlur(image, kernel_size, sigmaX) # 将图像缩放到 (32, 32, 3) resized_image = cv2.resize(blurred_image, (32, 32), interpolation=cv2.INTER_AREA) # 将处理后的图像从 numpy 数组转换回 tensor resized_image_tensor = torch.tensor(resized_image, dtype=torch.uint8, device=hdr_latents.device).permute(2, 0, 1) # 转回 (3, 32, 32) # 将处理后的图像添加到列表中 processed_images.append(resized_image_tensor) # 将列表中的所有图像堆叠成一个 tensor processed_images_tensor = torch.stack(processed_images) # 形状为 (16, 3, 32, 32) return processed_images_tensor def normalize_hdr(self, img): img = img / 255.0 return img * 2 -1 def forward(self, hdr_latents): # import pdb; pdb.set_trace() # todo: mask get hdr1, hdr2; input hdr_latents(实际上暂时是ldr) # 输入的形状为 (n, 16, 3, 256, 256)，去掉多余的维度 # import pdb; pdb.set_trace() # hdr_latents = hdr_latents.squeeze(0) # 变成 (16, 3, 256, 256) batch_size_ori = hdr_latents.shape[0] # frame_num = hdr_latents.shape[1] hdr_latents = rearrange(hdr_latents, 'b f c h w -> (b f) c h w') batch_size = hdr_latents.shape[0] # batch_size = hdr_latents.shape[0] # 转换为 NCHW 形式： (batch, channels, height, width) 输入之前numpy2tensor已经permute过了 # 高斯模糊 hdr_latents = self.blur_image(hdr_latents) #(16,3,32,32) 可视化打印一下! # import pdb; pdb.set_trace() # todo: decompose hdr hdr_latents_1, hdr_latents_2 = self.decompose_hdr(hdr_latents) # [16, 3, 32, 32], [16, 3, 32, 32] # todo: 加一步 normalize /255 -> -1,1 hdr_latents, hdr_latents_1, hdr_latents_2 = self.normalize_hdr(hdr_latents), self.normalize_hdr(hdr_latents_1), self.normalize_hdr(hdr_latents_2) # import pdb; pdb.set_trace() # 将输出展平，准备输入到 MLP 中 hdr_latents = hdr_latents.reshape(hdr_latents.size(0), -1) # [16, 3072] hdr_latents_1 = hdr_latents_1.reshape(hdr_latents_1.size(0), -1) # [16, 3072] hdr_latents_2 = hdr_latents_2.reshape(hdr_latents_2.size(0), -1) # 传递给 MLP hdr_latents = self.mlp(hdr_latents) #(16, 2304) 3072 -> 2304 hdr_latents_1 = self.mlp(hdr_latents_1) hdr_latents_2 = self.mlp(hdr_latents_2) # 重新调整输出的形状 hdr_latents = hdr_latents.reshape(batch_size, 3, 768) # reshape 输出为 (16*n, 3, 768) hdr_latents_1 = hdr_latents_1.reshape(batch_size, 3, 768) hdr_latents_2 = hdr_latents_2.reshape(batch_size, 3, 768) hdr_latents = rearrange(hdr_latents, '(b f) n c -> b f n c', b=batch_size_ori) hdr_latents_1 = rearrange(hdr_latents_1, '(b f) n c -> b f n c', b=batch_size_ori) hdr_latents_2 = rearrange(hdr_latents_2, '(b f) n c -> b f n c', b=batch_size_ori) #! 两个细节: 1. 仅有ldr, 需不需要concat hdr或线性变换 2. mask不同帧不一致 return hdr_latents, hdr_latents_1, hdr_latents_2 # 3 x (b,16,3,768) class InstructP2PVideoTrainer(DDIMLDMTextTraining): def __init__( self, *args, cond_image_dropout=0.1, cond_text_dropout=0.1, cond_hdr_dropout=0.1, prompt_type='output_prompt', text_cfg=7.5, img_cfg=1.2, hdr_cfg=7.5, hdr_rate=0.1, ic_condition='bg', hdr_train=False, **kwargs ): super().__init__(*args, **kwargs) self.hdr_train = hdr_train if self.hdr_train: self.hdr_encoder = HDRCtrlModel() self.hdr_encoder = self.hdr_encoder.to(self.unet.device) self.mlp = CombineMLP() self.cond_hdr_dropout = cond_hdr_dropout self.hdr_rate = hdr_rate self.cond_image_dropout = cond_image_dropout self.cond_text_dropout = cond_text_dropout assert ic_condition in ['fg', 'bg'] assert prompt_type in ['output_prompt', 'edit_prompt', 'mixed_prompt'] self.prompt_type = prompt_type self.ic_condition = ic_condition self.text_cfg = text_cfg self.img_cfg = img_cfg self.hdr_cfg = hdr_cfg #! 开启xformers训练设置 # self.unet.enable_xformers_memory_efficient_attention() # self.unet.enable_gradient_checkpointing() def encode_text(self, text): with torch.cuda.amp.autocast(dtype=torch.float16): encoded_text = super().encode_text(text) return encoded_text def encode_image_to_latent(self, image): # with torch.cuda.amp.autocast(dtype=torch.float16): latent = super().encode_image_to_latent(image) return latent # @torch.cuda.amp.autocast(dtype=torch.float16) @torch.no_grad() def get_prompt(self, batch, mode): # if mode == 'train': # if self.prompt_type == 'output_prompt': # prompt = batch['output_prompt'] # elif self.prompt_type == 'edit_prompt': # training的时候是edit prompt # prompt = batch['edit_prompt'] # elif self.prompt_type == 'mixed_prompt': # if int(torch.rand(1)) > 0.5: # prompt = batch['output_prompt'] # else: # prompt = batch['edit_prompt'] # else: # prompt = batch['output_prompt'] prompt = batch['text_prompt'] if not self.hdr_train: #! 如果hdr后续加进来text了, 还是需要? if torch.rand(1).item() < self.cond_text_dropout: prompt = 'change the background' cond_text = self.encode_text(prompt) if mode == 'train': if torch.rand(1).item() < self.cond_text_dropout: cond_text = torch.zeros_like(cond_text) # import pdb; pdb.set_trace() return cond_text # @torch.cuda.amp.autocast(dtype=torch.float16) @torch.no_grad() def encode_image_to_latent(self, image): b, f, c, h, w = image.shape image = rearrange(image, 'b f c h w -> (b f) c h w') latent = super().encode_image_to_latent(image) latent = rearrange(latent, '(b f) c h w -> b f c h w', b=b) return latent # @torch.cuda.amp.autocast(dtype=torch.float16) @torch.no_grad() def decode_latent_to_image(self, latent): b, f, c, h, w = latent.shape latent = rearrange(latent, 'b f c h w -> (b f) c h w') image = [] for latent_ in latent: image_ = super().decode_latent_to_image(latent_[None]) image.append(image_.sample) #! 注意一下这里 之前没报过错吗; -> 之前不是一个类 image = torch.cat(image, dim=0) # image = super().decode_latent_to_image(latent) image = rearrange(image, '(b f) c h w -> b f c h w', b=b) return image @torch.no_grad() def get_cond_image(self, batch, mode): # import pdb; pdb.set_trace() cond_fg_image = batch['fg_video'] # 这边condition 就是 input_video了, 估计是concat或者ctrlnet cond_fg_image = self.encode_image_to_latent(cond_fg_image) if self.ic_condition == 'bg': cond_bg_image = batch['bg_video'] if torch.all(cond_bg_image == 0): cond_bg_image = torch.zeros_like(cond_fg_image) #! 背景一定概率为0, 置为0.3 else: cond_bg_image = self.encode_image_to_latent(cond_bg_image) cond_image = torch.cat((cond_fg_image, cond_bg_image), dim=2) #(1,16,8,64,64) else: cond_image = cond_fg_image # test code: 可视化代码 # from PIL import Image # Image.fromarray(((batch['input_video'] + 1) / 2 * 255).byte()[0,0].permute(1,2,0).cpu().numpy()).save('img1.png') # ip2p does not scale cond image, so we unscale the cond image # cond_image = self.encode_image_to_latent(cond_image) / self.scale_factor # 额 就是一个vae encode，没有缩放；这边不进行缩放吗? 啥意思呢 if mode == 'train': # if int(torch.rand(1)) < self.cond_image_dropout: # 0.1的概率随机初始化, 应该是为了保障一个鲁棒性 难怪有的时候是全0, 不是代码的bug #! 艹 bug, 这么久才发现.... if torch.rand(1).item() < self.cond_image_dropout: cond_image = torch.zeros_like(cond_image) return cond_image @torch.no_grad() def get_diffused_image(self, batch, mode): # import pdb; pdb.set_trace() x = batch['tgt_video'] # 这边编辑的时候, 具体加噪和去噪的gt, 整个这套流程都是以编辑后, 即edited video作为输入 # from PIL import Image # Image.fromarray(((batch['edited_video'] + 1) / 2 * 255).byte()[0,0].permute(1,2,0).cpu().numpy()).save('img2.png') b, *_ = x.shape x = self.encode_image_to_latent(x) # (1, 16, 4, 32, 32), 经过了vae encode eps = torch.randn_like(x) if mode == 'train': t = torch.randint(0, self.num_timesteps, (b,), device=x.device).long() else: t = torch.full((b,), self.num_timesteps-1, device=x.device, dtype=torch.long) x_t = self.add_noise(x, t, eps) # 加噪t步长 eps表示高斯噪声, 和scheduler的加噪 if self.prediction_type == 'epsilon': return x_t, eps, t else: return x_t, x, t @torch.no_grad() def get_hdr_image(self, batch, mode): x = batch['ldr_video'] # todo (16,3,256,512), float, tensor, device -> (1,16,3,256,256) 注意此时仅有ldr # import pdb; pdb.set_trace() hdr_latents, hdr_latents_1, hdr_latents_2 = self.hdr_encoder(x) if mode == 'train': #! 考虑一下这个开不开, 因为后面要拉consistency loss if torch.rand(1).item() < self.cond_hdr_dropout: hdr_latents = torch.zeros_like(hdr_latents) hdr_latents_1 = torch.zeros_like(hdr_latents_1) hdr_latents_2 = torch.zeros_like(hdr_latents_2) return hdr_latents, hdr_latents_1, hdr_latents_2 @torch.no_grad() # batch中需要加载mask def get_mask(self, batch, mode, target): # (1,16,1,512,512) # import pdb; pdb.set_trace() mask = batch['fg_mask'] # todo 返回mask (n,16,1,512,512) bs = mask.shape[0] target_height, target_width = target.shape[-2:] #(n,16,3,64,64) mask = rearrange(mask, 'b f c h w -> (b f) c h w') resized_mask = F.interpolate(mask, size=(target_height, target_width), mode='bilinear', align_corners=False) # resized_mask = resized_mask.unsqueeze(0) resized_mask = rearrange(resized_mask, '(b f) c h w -> b f c h w', b=bs) if target.shape[2] != resized_mask.shape[2]: resized_mask = resized_mask.expand(-1, -1, target.shape[2], -1, -1) # 匹配目标通道数 return resized_mask @torch.no_grad() def process_batch(self, batch, mode): #! 可视化这边的image, 查看问题出在哪了。。。 √, 应该是randn_drop的事 # import pdb; pdb.set_trace() cond_image = self.get_cond_image(batch, mode) # 把输入的src image进行一个编码, 这边只有vae的encode, 且没有乘缩放的系数(ip2p本身没乘...) diffused_image, target, t = self.get_diffused_image(batch, mode) # diffused_image: 经过了vae encode, 和scheduler的加噪，标准的降噪输入 # target: 这边是epsilon目标, 因此还是拉成epsilon的损失；t: 训练阶段是随机的一个数值, 推理阶段一般都是1000 prompt = self.get_prompt(batch, mode) model_kwargs = { 'encoder_hidden_states': prompt } # import pdb; pdb.set_trace() if self.hdr_train: hdr_image, hdr_image_1, hdr_image_2 = self.get_hdr_image(batch, mode) #(16,3,768) fg_mask = self.get_mask(batch, mode, target) # 把原图像前景mask resize到target大小 model_kwargs = { 'encoder_hidden_states': {'hdr_latents': hdr_image, 'encoder_hidden_states': prompt, 'hdr_latents_1': hdr_image_1, 'hdr_latents_2': hdr_image_2, 'fg_mask': fg_mask} } return { 'diffused_input': diffused_image, # (1, 16, 4, 64, 64), 经过了vae encode, 和scheduler的加噪 'condition': cond_image, # 把输入的src image进行一个编码, 这边只有vae的encode, 且没有乘缩放的系数 (1,16,8,64,64) 'target': target, # 这个是加到tgt video的高斯噪声 't': t, # 0~1000的一个时刻 'model_kwargs': model_kwargs, # 这边就是一个text_hidden_states } def training_step(self, batch, batch_idx): #! 注意一下仅仅训motion layer # import pdb; pdb.set_trace() processed_batch = self.process_batch(batch, mode='train') #(1,16,3,256,256), 读取的序列化图片, 仅仅做了一个归一化操作 diffused_input = processed_batch['diffused_input'] # (1,16,4,64,64), edit images, 经过了vae encode, 和scheduler的加噪 condition = processed_batch['condition'] # (1,16,8,64,64) 把输入的src images进行一个编码, 这边只有vae的encode, 且没有乘缩放的系数 target = processed_batch['target'] # (1,16,4,64,64), target是加入的高斯噪声 t = processed_batch['t'] # [257], 一个0~1000的随机时刻 model_kwargs = processed_batch['model_kwargs'] # dict, 仅包含一项: encoder_hidden_states, [1, 77, 768] text_hidden_states model_input = torch.cat([diffused_input, condition], dim=2) # b, f, c, h, w [1,16,8,32,32] 这边是做的concat, 很多edit文章经典操作, 把两个东西concat起来 #! 半精度 # model_input = model_input.float() # model_kwargs['encoder_hidden_states'] = model_kwargs['encoder_hidden_states'].half() model_input = rearrange(model_input, 'b f c h w -> b c f h w') # [1,8,16,32,32] pred = self.unet(model_input, t, **model_kwargs).sample # (1,4,16,64,64) #! pred = rearrange(pred, 'b c f h w -> b f c h w') # (1,16,4,64,64) #! if not self.hdr_train: loss = self.get_loss(pred, target, t) # 0.320 else: fg_mask = model_kwargs['encoder_hidden_states']['fg_mask'] loss = self.get_hdr_loss(fg_mask, pred, target) ### add consistency loss ### # todo: 三个相同的model_input, 不同的model_kwargs (注意stack到一起, attn里面的逻辑也得改...) # if self.hdr_train: # fg_mask = model_kwargs['encoder_hidden_states']['fg_mask'] # hdr_latents = model_kwargs['encoder_hidden_states']['hdr_latents_1'] # hdr_latents_1 = model_kwargs['encoder_hidden_states']['hdr_latents_1'] # hdr_latents_2 = model_kwargs['encoder_hidden_states']['hdr_latents_2'] # model_input = torch.cat([diffused_input, condition], dim=2) # model_input = rearrange(model_input, 'b f c h w -> b c f h w') # model_input_1 = model_input.clone() # model_input_2 = model_input.clone() # model_input_all = torch.cat([model_input, model_input_1, model_input_2], dim=0) # prompt = model_kwargs['encoder_hidden_states']['encoder_hidden_states'] #(1*n,77,768) # prompt_all = torch.cat([prompt, prompt, prompt], dim=0) #(3*n,77,768) # # import pdb; pdb.set_trace() # model_kwargs['encoder_hidden_states']['encoder_hidden_states'] = prompt_all # # import pdb; pdb.set_trace() # hdr_latents_all = torch.cat([hdr_latents, hdr_latents_1, hdr_latents_2], dim=0) #(3*n,16,77,768) # model_kwargs['encoder_hidden_states']['hdr_latents']=hdr_latents_all # pred_all = self.unet(model_input_all, t, **model_kwargs).sample # (1,4,16,64,64) # pred_all = rearrange(pred_all, 'b c f h w -> b f c h w') # pred, pred1, pred2 = pred_all.chunk(3, dim=0) # loss_ori = self.get_hdr_loss(fg_mask, pred, target) # # 假设获得了L1, L2 # # hdr_latents_1 = mask(hdr_latents) # 随机构造一个mask + 逻辑矫正 # # model_kwargs['encoder_hidden_states']['hdr_latents']=hdr_latents_1 # # pred1 = self.unet(model_input, t, **model_kwargs).sample # get L1下的预测值 (1,16,4,64,64) # # pred1 = rearrange(pred1, 'b c f h w -> b f c h w') # # model_input = torch.cat([diffused_input, condition], dim=2) # # model_input = rearrange(model_input, 'b f c h w -> b c f h w') # # # hdr_latents_2 = 1-mask(hdr_latents) # # model_kwargs['encoder_hidden_states']['hdr_latents']=hdr_latents_2 # # pred2 = self.unet(model_input, t, **model_kwargs).sample # get L2下的预测值 # # pred2 = rearrange(pred2, 'b c f h w -> b f c h w') # # import pdb; pdb.set_trace() # pred_combine = self.mlp(pred1, pred2) #! todo: 构造mlp loss 错了!! 搞对一下, 应该需要展平.... # loss_c = self.get_hdr_loss(fg_mask, pred, pred_combine) # # loss_c = MSELoss(mask*pred, mask*pred_conbine) # todo: change to函数, 逻辑矫正 # loss = loss_ori + self.hdr_rate * loss_c # 设一个系数, 好控制变化 ### end ### self.log('train_loss', loss, sync_dist=True) latent_pred = self.predict_x_0_from_x_t(pred, t, diffused_input) # (1,16,4,32,32) image_pred = self.decode_latent_to_image(latent_pred) # 这边相当于是pred_x0了, (1,16,3,256,256) drop_out = torch.all(condition == 0).item() res_dict = { 'loss': loss, 'pred': image_pred, 'drop_out': drop_out, 'time': t[0].item() } return res_dict @torch.no_grad() @torch.cuda.amp.autocast(dtype=torch.bfloat16) def validation_step(self, batch, batch_idx): # 没写好 可以先pass # pass # import pdb; pdb.set_trace() if not self.hdr_train: from .inference.inference import InferenceIP2PVideo inf_pipe = InferenceIP2PVideo( self.unet, beta_start=self.scheduler.config.beta_start, beta_end=self.scheduler.config.beta_end, beta_schedule=self.scheduler.config.beta_schedule, num_ddim_steps=20 ) # import pdb; pdb.set_trace() processed_batch = self.process_batch(batch, mode='val') diffused_input = torch.randn_like(processed_batch['diffused_input']) #(1,16,4,64,64) condition = processed_batch['condition'] # 这边其实留有一个接口给condition (1,16,8,64,64) img_cond = condition text_cond = processed_batch['model_kwargs']['encoder_hidden_states'] # import pdb; pdb.set_trace() res = inf_pipe( latent = diffused_input, text_cond = text_cond, text_uncond = self.encode_text(['']), img_cond = img_cond, text_cfg = self.text_cfg, img_cfg = self.img_cfg, hdr_cfg = self.hdr_cfg ) latent_pred = res['latent'] image_pred = self.decode_latent_to_image(latent_pred) res_dict = { 'pred': image_pred, } else: from .inference.inference import InferenceIP2PVideoHDR inf_pipe = InferenceIP2PVideoHDR( self.unet, beta_start=self.scheduler.config.beta_start, beta_end=self.scheduler.config.beta_end, beta_schedule=self.scheduler.config.beta_schedule, num_ddim_steps=20 ) # import pdb; pdb.set_trace() processed_batch = self.process_batch(batch, mode='val') diffused_input = torch.randn_like(processed_batch['diffused_input']) #(1,16,4,64,64) condition = processed_batch['condition'] # 这边其实留有一个接口给condition (1,16,8,64,64) model_kwargs = processed_batch['model_kwargs'] img_cond = condition text_cond = model_kwargs['encoder_hidden_states']['encoder_hidden_states'] hdr_cond = model_kwargs['encoder_hidden_states']['hdr_latents'] # import pdb; pdb.set_trace() res = inf_pipe( latent = diffused_input, text_cond = text_cond, text_uncond = self.encode_text(['']), hdr_cond = hdr_cond, img_cond = img_cond, text_cfg = self.text_cfg, img_cfg = self.img_cfg, ) latent_pred = res['latent'] image_pred = self.decode_latent_to_image(latent_pred) res_dict = { 'pred': image_pred, } return res_dict def configure_optimizers(self): # optimizer = torch.optim.AdamW(self.unet.parameters(), lr=self.optim_args['lr']) import bitsandbytes as bnb params = [] for name, p in self.unet.named_parameters(): if ('transformer_in' in name) or ('temp_' in name): # p.requires_grad = True params.append(p) else: pass # p.requires_grad = False optimizer = bnb.optim.Adam8bit(params, lr=self.optim_args['lr'], betas=(0.9, 0.999)) return optimizer def initialize_unet(self, unet_init_weights): if unet_init_weights is not None: print(f'INFO: initialize denoising UNet from {unet_init_weights}') sd = torch.load(unet_init_weights, map_location='cpu') model_sd = self.unet.state_dict() # fit input conv size for k in model_sd.keys(): if k in sd.keys(): pass else: # handling temporal layers if (('temp_' in k) or ('transformer_in' in k)) and 'proj_out' in k: # print(f'INFO: initialize {k} from {model_sd[k].shape} to zeros') sd[k] = torch.zeros_like(model_sd[k]) else: # print(f'INFO: initialize {k} from {model_sd[k].shape} to random') sd[k] = model_sd[k] self.unet.load_state_dict(sd) class InstructP2PVideoTrainerTemporal(InstructP2PVideoTrainer): def initialize_unet(self, unet_init_weights): # 这边对比上一级来说, 新加的部分在于 rewrite了unet的load函数 if unet_init_weights is not None: print(f'INFO: initialize denoising UNet from {unet_init_weights}') sd_init_weights, motion_module_init_weights, iclight_init_weights = unet_init_weights os.makedirs(sd_init_weights, exist_ok=True) sd_init_weights, motion_module_init_weights, iclight_init_weights = f'models/{sd_init_weights}', f'models/{motion_module_init_weights}', f'models/{iclight_init_weights}' if not os.path.exists(sd_init_weights): url = 'https://huggingface.co/stablediffusionapi/realistic-vision-v51/resolve/main/unet/diffusion_pytorch_model.safetensors' download_url_to_file(url=url, dst=sd_init_weights) if not os.path.exists(motion_module_init_weights): url = 'https://huggingface.co/aleafy/RelightVid/resolve/main/relvid_mm_sd15_fbc.pth' download_url_to_file(url=url, dst=motion_module_init_weights) if not os.path.exists(iclight_init_weights): url = 'https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fbc.safetensors' download_url_to_file(url=url, dst=iclight_init_weights) sd = load_file(sd_init_weights) #! 关于加载iclight的unet, 后面再加到yaml里面... 我甚至觉得只要改unet, vae和text其实都差不多 # sd = torch.load(sd_init_weights, map_location='cpu') # 注意debug看看这是啥 + 打印一下原有和加载的keys if self.unet.use_motion_module: motion_sd = torch.load(motion_module_init_weights, map_location='cpu') assert len(sd) + len(motion_sd) == len(self.unet.state_dict()), f'Improper state dict length, got {len(sd) + len(motion_sd)} expected {len(self.unet.state_dict())}' #! 注意一下这行保证了加载的key至少在数量上是对应的; 这行的目的是self.unet是自己定义的 而这两个加载的是别的地方训练的(可能是diffusers中的) sd.update(motion_sd) for k, v in self.unet.state_dict().items(): if 'pos_encoder.pe' in k: # 这边是原来iv2v的代码 temporal_position_encoding_max_len, 设置为 32 sd[k] = v # the size of pe may change, 主要是temporal layer的size会发生改变... √ 由于输入的max_len变了 # if 'conv_in.weight' in k: #! tmp, 这里是test一下 # sd[k] = v else: assert len(sd) == len(self.unet.state_dict()) self.unet.load_state_dict(sd) # 为什么这里可以完美适配? √ # todo: 更改sd的conv_in.weight的shape到12; 更改函数forward, 支持多个输入cond; iclight的sd_offset加载进去; unet = self.unet # saVe一下 # 这里是更改conv_in的shape; #! 这边注意一下要改成3D版本的unet with torch.no_grad(): # new_conv_in = torch.nn.Conv2d(12, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding) new_conv_in = InflatedConv3d(12, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding) new_conv_in.weight.zero_() new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight) new_conv_in.bias = unet.conv_in.bias unet.conv_in = new_conv_in ###### -- 更改 forward函数 --- ##### # 这里是更改forward函数。 具体调用的部分在main后面，那里也得改 # unet_original_forward = unet.forward # def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs): # c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample) # (1,8,67,120) # c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0) # (2,8,67,120) 应该是复制一份，用于cfg # new_sample = torch.cat([sample, c_concat], dim=2) #(2,12,67,120) 这边还是在通道维度上进行的concat #! change 在第二维cat (2,1,12,67,120) # # todo 这边中间可以加一个f的通道 b,c,f,h,w ; 另一种方式: 对于数据进行改变, 那么上述concat的代码也需要进行变换了... # # new_sample = new_sample.unsqueeze(2) # (2,12,1,67,120) #! 这里需要change, 要在一输入之前就要更改他的维度, 因此前面concat也需要稍微改一下 不要在forward中增加f维度 (因为要依赖输入) # new_sample = rearrange(new_sample, 'b f c h w -> b c f h w') # kwargs['cross_attention_kwargs'] = {} # # return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs) # result = unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs) # # return (result[0].squeeze(2),) #! tmp # return (rearrange(result[0], 'b c f h w -> b f c h w'),) # unet.forward = hooked_unet_forward ##### -- 更改 forward函数 --- ##### # model_path = '/home/fy/Code/instruct-video-to-video/IC-Light/models/iclight_sd15_fbc.safetensors' # 这里是加载iclight的lora weight sd_offset = load_file(iclight_init_weights) sd_origin = unet.state_dict() keys = sd_origin.keys() for k in sd_offset.keys(): sd_origin[k] = sd_origin[k] + sd_offset[k] # sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()} self.unet.load_state_dict(sd_origin, strict=True) del sd_offset, sd_origin, unet, keys # print(1) # todo 试写一下iclight unet的加载方式 # sd = load_file('/home/fy/Code/IC-Light/cache_models/models--stablediffusionapi--realistic-vision-v51/snapshots/19e3643d7d963c156d01537188ec08f0b79a514a/unet/diffusion_pytorch_model.safetensors') # debug: print参数 # with open('logs/sd_keys.txt', 'w') as f: # f.write("SD Keys:\n") # for key in sd_ori.keys(): # f.write(f"{key}\n") # unet_state_dict = self.unet.state_dict() # with open('logs/unet_state_dict_keys.txt', 'w') as f: # f.write("UNet State Dict Keys:\n") # for key in unet_state_dict.keys(): # f.write(f"{key}\n") else: with torch.no_grad(): new_conv_in = InflatedConv3d(12, self.unet.conv_in.out_channels, self.unet.conv_in.kernel_size, self.unet.conv_in.stride, self.unet.conv_in.padding) self.unet.conv_in = new_conv_in def configure_optimizers(self): # 决定了仅仅训练motion_module的参数 注意一下pl.Trainer独有的函数 import bitsandbytes as bnb motion_params = [] remaining_params = [] train_names = [] # for debug for name, p in self.unet.named_parameters(): if ('motion' in name): #! 哦哦 这里决定了哪些参数用于训练... 这里实际训练的只有motion相关参数 motion_params.append(p) train_names.append(name) elif ('attentions' in name): motion_params.append(p) train_names.append(name) else: remaining_params.append(p) # import pdb; pdb.set_trace() optimizer = bnb.optim.Adam8bit([ {'params': motion_params, 'lr': self.optim_args['lr']}, ], betas=(0.9, 0.999)) return optimizer class InstructP2PVideoTrainerTemporalText(InstructP2PVideoTrainerTemporal): def initialize_unet(self, unet_init_weights): # 这边对比上一级来说, 新加的部分在于 rewrite了unet的load函数 if unet_init_weights is not None: print(f'INFO: initialize denoising UNet from {unet_init_weights}') sd_init_weights, motion_module_init_weights, iclight_init_weights = unet_init_weights if self.base_path: sd_init_weights = f"{self.base_path}/{sd_init_weights}" if '.safetensors' in sd_init_weights: # .safetensors的加载方式 sd = load_file(sd_init_weights) #! 关于加载iclight的unet, 后面再加到yaml里面... 我甚至觉得只要改unet, vae和text其实都差不多 else: #'.ckpt'场景 sd = torch.load(sd_init_weights, map_location='cpu') # sd = torch.load(sd_init_weights, map_location='cpu') # 注意debug看看这是啥 + 打印一下原有和加载的keys if self.unet.use_motion_module: motion_sd = torch.load(motion_module_init_weights, map_location='cpu') assert len(sd) + len(motion_sd) == len(self.unet.state_dict()), f'Improper state dict length, got {len(sd) + len(motion_sd)} expected {len(self.unet.state_dict())}' #! 注意一下这行保证了加载的key至少在数量上是对应的; 这行的目的是self.unet是自己定义的 而这两个加载的是别的地方训练的(可能是diffusers中的) sd.update(motion_sd) for k, v in self.unet.state_dict().items(): if 'pos_encoder.pe' in k: # 这边是原来iv2v的代码 temporal_position_encoding_max_len, 设置为 32 sd[k] = v # the size of pe may change, 主要是temporal layer的size会发生改变... √ 由于输入的max_len变了 # if 'conv_in.weight' in k: #! tmp, 这里是test一下 # sd[k] = v else: assert len(sd) == len(self.unet.state_dict()) self.unet.load_state_dict(sd) # 为什么这里可以完美适配? √ # todo: 更改sd的conv_in.weight的shape到12; 更改函数forward, 支持多个输入cond; iclight的sd_offset加载进去; unet = self.unet # saVe一下 # 这里是更改conv_in的shape; #! 这边注意一下要改成3D版本的unet with torch.no_grad(): # new_conv_in = torch.nn.Conv2d(12, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding) new_conv_in = InflatedConv3d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding) new_conv_in.weight.zero_() new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight) new_conv_in.bias = unet.conv_in.bias unet.conv_in = new_conv_in # model_path = '/home/fy/Code/instruct-video-to-video/IC-Light/models/iclight_sd15_fbc.safetensors' # 这里是加载iclight的lora weight sd_offset = load_file(iclight_init_weights) sd_origin = unet.state_dict() keys = sd_origin.keys() for k in sd_offset.keys(): sd_origin[k] = sd_origin[k] + sd_offset[k] # sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()} self.unet.load_state_dict(sd_origin, strict=True) del sd_offset, sd_origin, unet, keys else: with torch.no_grad(): new_conv_in = InflatedConv3d(8, self.unet.conv_in.out_channels, self.unet.conv_in.kernel_size, self.unet.conv_in.stride, self.unet.conv_in.padding) self.unet.conv_in = new_conv_in