From a6143c8858177ac7461af1be3bef30c44646b822 Mon Sep 17 00:00:00 2001 From: FuouM Date: Mon, 1 Jul 2024 14:42:25 +0700 Subject: [PATCH] face parser support --- .gitignore | 2 + README.md | 2 + constants.py | 14 +- face_parsing/face_parsing_loader.py | 25 +++ face_parsing/face_parsing_model.py | 275 ++++++++++++++++++++++++++++ face_parsing/resnet.py | 97 ++++++++++ partswap_inference.py | 31 +++- run.py | 22 ++- 8 files changed, 459 insertions(+), 9 deletions(-) create mode 100644 face_parsing/face_parsing_loader.py create mode 100644 face_parsing/face_parsing_model.py create mode 100644 face_parsing/resnet.py diff --git a/.gitignore b/.gitignore index b233313..7e54636 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,8 @@ checkpoints/* !checkpoints/place_checkpoints_here.txt +face_parsing/*.pth +!face_parsing/put_model_safetensors_config_json_here.txt # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index 4b2648f..7f44abd 100644 --- a/README.md +++ b/README.md @@ -55,3 +55,5 @@ place_checkpoints_here.txt vox-adv-cpk.pth.tar vox-cpk.pth.tar ``` + +resnet18_url = "https://download.pytorch.org/models/resnet18-5c106cde.pth" diff --git a/constants.py b/constants.py index 9373269..2eade3a 100644 --- a/constants.py +++ b/constants.py @@ -21,22 +21,28 @@ config_suffix = "256.yaml" config_folder = "config" checkpoint_folder = "checkpoints" +face_parser_checkpoint_name = "79999_iter.pth" + partswap_model_config_dict = { "vox-5segments": "vox-256-sem-5segments.yaml", "vox-10segments": "vox-256-sem-10segments.yaml", "vox-15segments": "vox-256-sem-15segments.yaml", "vox-cpk": "vox-256-sem-10segments.yaml", "vox-adv-cpk": "vox-256-sem-10segments.yaml", + "vox-first-order": "vox-256-sem-10segments.yaml", } -partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk"] +partswap_face_parser_enforce_models = ["vox-first-order"] + +partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk", "vox-first-order"] partswap_model_length_dict = { "vox-5segments": 6, "vox-10segments": 11, - "vox-15segments": 15, - "vox-cpk": 19, - "vox-adv-cpk": 19, + "vox-15segments": 16, + "vox-cpk": 11, + "vox-adv-cpk": 11, + "vox-first-order": 11, } partswap_model_names = list(partswap_model_config_dict.keys()) diff --git a/face_parsing/face_parsing_loader.py b/face_parsing/face_parsing_loader.py new file mode 100644 index 0000000..7b8118c --- /dev/null +++ b/face_parsing/face_parsing_loader.py @@ -0,0 +1,25 @@ +import torch +import numpy as np + +from .face_parsing_model import BiSeNet + +def load_face_parser_model(base_dir: str, face_parser_checkpoint_name: str, cpu=False): + print("Loading face_parser model") + face_parser = BiSeNet(n_classes=19) + face_parser_path = f"{base_dir}/face_parsing/{face_parser_checkpoint_name}" + + mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view(1, 3, 1, 1) + std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view(1, 3, 1, 1) + + if not cpu: + face_parser.cuda() + face_parser.load_state_dict(torch.load(face_parser_path)) + face_parser.mean = mean.cuda() + face_parser.std = std.cuda() + else: + face_parser.load_state_dict(torch.load(face_parser_path, map_location=torch.device('cpu'))) + face_parser.mean = mean + face_parser.std = std + print("face_parser model loaded!") + + return face_parser \ No newline at end of file diff --git a/face_parsing/face_parsing_model.py b/face_parsing/face_parsing_model.py new file mode 100644 index 0000000..718ea14 --- /dev/null +++ b/face_parsing/face_parsing_model.py @@ -0,0 +1,275 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .resnet import Resnet18 + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d( + in_chan, + out_chan, + kernel_size=ks, + stride=stride, + padding=padding, + bias=False, + ) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode="nearest") + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode="nearest") + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode="nearest") + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d( + out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False + ) + self.conv2 = nn.Conv2d( + out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False + ) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode="bilinear", align_corners=True) + feat_out16 = F.interpolate( + feat_out16, (H, W), mode="bilinear", align_corners=True + ) + feat_out32 = F.interpolate( + feat_out32, (H, W), mode="bilinear", align_corners=True + ) + return feat_out, feat_out16, feat_out32 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance( + child, BiSeNetOutput + ): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params diff --git a/face_parsing/resnet.py b/face_parsing/resnet.py new file mode 100644 index 0000000..bd30bfd --- /dev/null +++ b/face_parsing/resnet.py @@ -0,0 +1,97 @@ +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.model_zoo as modelzoo +import torch +from pathlib import Path + + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum - 1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + base_dir = Path(__file__).resolve().parent + print(f"resnet base dir {base_dir}") + state_dict = torch.load(f"{base_dir}/resnet18-5c106cde.pth") + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if "fc" in k: + continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params diff --git a/partswap_inference.py b/partswap_inference.py index 79a9859..15a5c2c 100644 --- a/partswap_inference.py +++ b/partswap_inference.py @@ -1,5 +1,6 @@ import numpy as np import torch +import torch.nn.functional as F from tqdm import tqdm from comfy.utils import ProgressBar @@ -7,6 +8,15 @@ from .modules.segmentation_module import SegmentationModule from .partswap_model import PartSwapGenerator +def face_parse_seg( + source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64) +): + seg = F.interpolate(source_image, size=in_size) + seg = (seg - face_parser_model.mean) / face_parser_model.std + seg = torch.softmax(face_parser_model(seg)[0], dim=1) + seg = F.interpolate(seg, size=out_size) + return seg + def partswap_inference( swap_indices: list[int], source_image, @@ -14,6 +24,7 @@ def partswap_inference( reconstruction_module: PartSwapGenerator, segmentation_module: SegmentationModule, use_source_seg: bool, + face_parser_model=None, hard_edges=False, cpu=False, ) -> list: @@ -26,6 +37,7 @@ def partswap_inference( print("Using GPU") seg_source = segmentation_module(source_image) + print(f"{seg_source['segmentation'].shape}") num_frames = target_video.shape[2] pbar = ProgressBar(num_frames) @@ -33,15 +45,26 @@ def partswap_inference( target_frame = target_video[:, :, frame_idx] seg_target = segmentation_module(target_frame) seg_targets.append(seg_target) - if use_source_seg: - blend_mask = seg_source["segmentation"] + + if face_parser_model is not None: + blend_mask = face_parse_seg( + source_image if use_source_seg else target_frame, + face_parser_model, + (512, 512), + (64, 64), + ) + else: - blend_mask = seg_target["segmentation"] + blend_mask = ( + seg_source["segmentation"] + if use_source_seg + else seg_target["segmentation"] + ) blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True) if hard_edges: blend_mask = (blend_mask > 0.5).type(blend_mask.type()) - + print(f"{blend_mask.shape=}") out = reconstruction_module( source_image, target_frame, diff --git a/run.py b/run.py index 951090c..3775eb5 100644 --- a/run.py +++ b/run.py @@ -18,6 +18,8 @@ from .constants import ( partswap_model_length_dict, partswap_model_names, supported_model_names, + face_parser_checkpoint_name, + partswap_face_parser_enforce_models ) from .fomm_inference import inference, inference_best_frame from .fomm_loader import load_checkpoint @@ -32,6 +34,7 @@ from .utils import ( out_video, reshape_image, ) +from .face_parsing.face_parsing_loader import load_face_parser_model base_dir = Path(__file__).resolve().parent @@ -219,6 +222,10 @@ class FOMM_Partswap: "BOOLEAN", {"default": False}, ), + "use_face_parser": ( + "BOOLEAN", + {"default": False}, + ), "chosen_seg_indices": ("STRING", {}), "viz_alpha": ( "FLOAT", @@ -254,6 +261,7 @@ class FOMM_Partswap: blend_scale: float, use_source_seg: bool, hard_edges: bool, + use_face_parser: bool, chosen_seg_indices: str, viz_alpha: float, audio=None, @@ -281,9 +289,18 @@ class FOMM_Partswap: reconstruction_module, segmentation_module = load_partswap_checkpoint( config_path, checkpoint_path, blend_scale, use_fomm=use_fomm ) + + face_parser_model = None + if model_name in partswap_face_parser_enforce_models: + assert use_face_parser, "Special supervised model, requires face_parser" + if use_face_parser: + face_parser_model = load_face_parser_model(base_dir, face_parser_checkpoint_name) + face_parser_model.cuda() + face_parser_model.eval() + print(f"Chosen {chosen_seg_indices}") swap_indices = deserialize_integers(chosen_seg_indices) - swap_indices = swap_indices[: partswap_model_length_dict[model_name]] + swap_indices = [x for x in swap_indices if x < partswap_model_length_dict[model_name]] print(f"Swapping {swap_indices}") params = { @@ -293,6 +310,7 @@ class FOMM_Partswap: "reconstruction_module": reconstruction_module, "segmentation_module": segmentation_module, "use_source_seg": use_source_seg, + "face_parser_model": face_parser_model, "hard_edges": hard_edges, } @@ -305,6 +323,8 @@ class FOMM_Partswap: output_images = out_video(predictions) + del face_parser_model + return ( seg_src_viz, seg_tgt_viz,