Kerry
4b731ff98b
Float project applicable to ComfyUI.Generates speaking portrait video frames from an image and audio.适用于ComfyUI的强大[FLOAT]项目，实现由肖像生成音频驱动的说话视频。
25 lines
842 B
Python
25 lines
842 B
Python
from torch import nn
|
|
from .encoder import Encoder
|
|
from .styledecoder import Synthesis
|
|
|
|
from models import BaseModel
|
|
|
|
class Generator(BaseModel):
|
|
def __init__(self, size, style_dim=512, motion_dim=20, channel_multiplier=1, blur_kernel=[1, 3, 3, 1]):
|
|
super().__init__()
|
|
|
|
self.enc = Encoder(size, style_dim, motion_dim)
|
|
self.dec = Synthesis(size, style_dim, motion_dim, blur_kernel, channel_multiplier)
|
|
|
|
def get_direction(self):
|
|
return self.dec.direction(None)
|
|
|
|
def synthesis(self, wa, alpha, feat):
|
|
img, flow = self.dec(wa, alpha, feat)
|
|
return img
|
|
|
|
def forward(self, img_source, img_drive, h_start=None):
|
|
wa, alpha, feats = self.enc(img_source, img_drive, h_start)
|
|
img_recon, flow = self.dec(wa, alpha, feats)
|
|
return {'d_hat': img_recon, 'flow': flow}
|
|
|