The Net class is the core neural network model that combines several components to animate a face in a static image. It integrates a reference UNet, a denoising UNet, a face locator, and projection models for image and audio.
Initialization Arguments:
reference_unet (UNet2DConditionModel): The reference UNet used for face animation.denoising_unet (UNet3DConditionModel): The denoising UNet used for face animation.face_locator (FaceLocator): The face locator model.reference_control_writer: Component for writing reference control.reference_control_reader: Component for reading reference control.imageproj: Image projection model.audioproj: Audio projection model.
Forward Pass Inputs:
noisy_latents (torch.Tensor): The noisy latents tensor.timesteps (torch.Tensor): The timesteps tensor.ref_image_latents (torch.Tensor): The reference image latents tensor.face_emb (torch.Tensor): The face embeddings tensor.audio_emb (torch.Tensor): The audio embeddings tensor.mask (torch.Tensor): Hard face mask for face locator.full_mask (torch.Tensor): Pose Mask.face_mask (torch.Tensor): Face Mask.lip_mask (torch.Tensor): Lip Mask.uncond_img_fwd (bool, optional): Flag for reference image unconditional forward pass.uncond_audio_fwd (bool, optional): Flag for audio unconditional forward pass.
class Net(nn.Module):
def __init__(
self,
reference_unet: UNet2DConditionModel,
denoising_unet: UNet3DConditionModel,
face_locator: FaceLocator,
reference_control_writer,
reference_control_reader,
imageproj,
audioproj,
):
# ... implementation
def forward(
self,
noisy_latents: torch.Tensor,
timesteps: torch.Tensor,
ref_image_latents: torch.Tensor,
face_emb: torch.Tensor,
audio_emb: torch.Tensor,
mask: torch.Tensor,
full_mask: torch.Tensor,
face_mask: torch.Tensor,
lip_mask: torch.Tensor,
uncond_img_fwd: bool = False,
uncond_audio_fwd: bool = False,
):
# ... implementation
