diff --git a/fantasyportrait/model.py b/fantasyportrait/model.py index 2ab6ca4..f61b2d8 100644 --- a/fantasyportrait/model.py +++ b/fantasyportrait/model.py @@ -180,7 +180,7 @@ class PortraitAdapter(nn.Module): ) return proj_model - def get_adapter_proj(self, adapter_fea=None): + def get_adapter_proj(self, adapter_fea=None, adapter_scale=1.0, mouth_scale=1.0, emo_scale=1.0): split_sizes = [6, 6, 30, 512] headpose, eye, emo, mouth = torch.split( adapter_fea, split_sizes, dim=-1 @@ -189,13 +189,13 @@ class PortraitAdapter(nn.Module): mouth = mouth.view(B * frames, 1, 512) emo = emo.view(B * frames, 1, 30) - mouth_fea = self.mouth_proj_model(mouth) - emo_fea = self.emo_proj_model(emo) + mouth_fea = self.mouth_proj_model(mouth) * mouth_scale + emo_fea = self.emo_proj_model(emo) * emo_scale mouth_fea = mouth_fea.view(B, frames, 16, 2048) emo_fea = emo_fea.view(B, frames, 4, 2048) - adapter_fea = self.proj_model(adapter_fea) + adapter_fea = self.proj_model(adapter_fea) * adapter_scale adapter_fea = adapter_fea.view(B, frames, 4, 2048) diff --git a/fantasyportrait/nodes.py b/fantasyportrait/nodes.py index 5f86d97..dbe6aac 100644 --- a/fantasyportrait/nodes.py +++ b/fantasyportrait/nodes.py @@ -73,6 +73,9 @@ class FantasyPortraitFaceDetector: "required": { "portrait_model": ("FANTASYPORTRAITMODEL",), "images": ("IMAGE",), + "adapter_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Scale for the adapter projection"}), + "mouth_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Scale for the mouth projection"}), + "emo_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Scale for the emotion projection"}), }, } @@ -81,7 +84,7 @@ class FantasyPortraitFaceDetector: FUNCTION = "detect" CATEGORY = "WanVideoWrapper" - def detect(self, images, portrait_model): + def detect(self, images, portrait_model, adapter_scale=1.0, mouth_scale=1.0, emo_scale=1.0): B, H, W, C = images.shape num_frames = ((B - 1) // 4) * 4 + 1 images = images.clone()[:num_frames] @@ -117,7 +120,7 @@ class FantasyPortraitFaceDetector: portrait_model = portrait_model["proj_model"] portrait_model.to(device) - adapter_proj = portrait_model.get_adapter_proj(head_emo_feat_all.to(device, dtype=portrait_model.dtype)) + adapter_proj = portrait_model.get_adapter_proj(head_emo_feat_all.to(device, dtype=portrait_model.dtype), adapter_scale=adapter_scale, mouth_scale=mouth_scale, emo_scale=emo_scale) portrait_model.to(offload_device) pos_idx_range = portrait_model.split_audio_adapter_sequence(adapter_proj.size(1), num_frames=num_frames)