Delete code

code/LIA_Model.py

@@ -1,39 +0,0 @@
-import torch
-import torch.nn as nn
-from networks.encoder import Encoder
-from networks.styledecoder import Synthesis
-
-# This part is modified from: https://github.com/wyhsirius/LIA
-class LIA_Model(torch.nn.Module):
-    def __init__(self, size = 256, style_dim = 512, motion_dim = 20, channel_multiplier=1, blur_kernel=[1, 3, 3, 1], fusion_type=''):
-        super().__init__()
-        self.enc = Encoder(size, style_dim, motion_dim, fusion_type)
-        self.dec = Synthesis(size, style_dim, motion_dim, blur_kernel, channel_multiplier)
-     
-    def get_start_direction_code(self, x_start, x_target, x_face, x_aug):
-        enc_dic = self.enc(x_start, x_target,  x_face, x_aug)
-        
-        wa, alpha, feats = enc_dic['h_source'], enc_dic['h_motion'], enc_dic['feats']
-        
-        return wa, alpha, feats
-    
-    def render(self, start, direction, feats):
-        return self.dec(start, direction, feats)
-    
-    def load_lightning_model(self, lia_pretrained_model_path):
-        selfState = self.state_dict()
-
-        state = torch.load(lia_pretrained_model_path, map_location='cpu')
-        for name, param in state.items():
-            origName = name;
-            
-            if name not in selfState:
-                name = name.replace("lia.", "")
-                if name not in selfState:
-                    print("%s is not in the model."%origName)
-                    # You can ignore those errors as some parameters are only used for training
-                    continue
-            if selfState[name].size() != state[origName].size():
-                print("Wrong parameter length: %s, model: %s, loaded: %s"%(origName, selfState[name].size(), state[origName].size()))
-                continue
-            selfState[name].copy_(param)

code/__init__.py

@@ -1 +0,0 @@
-from.LIA_Model import *

code/app.py

@@ -1,403 +0,0 @@
-
-import argparse
-from datetime import datetime
-from pathlib import Path
-import numpy as np
-import torch
-from PIL import Image
-import gradio as gr
-import shutil
-import librosa
-import python_speech_features
-import time
-from LIA_Model import LIA_Model
-import os
-from tqdm import tqdm
-import argparse
-import numpy as np
-from torchvision import transforms
-from templates import *
-import argparse
-import shutil
-from moviepy.editor import *
-import librosa
-import python_speech_features
-import importlib.util
-import time
-import os
-import time
-import numpy as np
-
-
-
-# Disable Gradio analytics to avoid network-related issues
-gr.analytics_enabled = False
-
-
-def check_package_installed(package_name):
-    package_spec = importlib.util.find_spec(package_name)
-    if package_spec is None:
-        print(f"{package_name} is not installed.")
-        return False
-    else:
-        print(f"{package_name} is installed.")
-        return True
-
-def frames_to_video(input_path, audio_path, output_path, fps=25):
-    image_files = [os.path.join(input_path, img) for img in sorted(os.listdir(input_path))]
-    clips = [ImageClip(m).set_duration(1/fps) for m in image_files]
-    video = concatenate_videoclips(clips, method="compose")
-
-    audio = AudioFileClip(audio_path)
-    final_video = video.set_audio(audio)
-    final_video.write_videofile(output_path, fps=fps, codec='libx264', audio_codec='aac')
-
-def load_image(filename, size):
-    img = Image.open(filename).convert('RGB')
-    img = img.resize((size, size))
-    img = np.asarray(img)
-    img = np.transpose(img, (2, 0, 1))  # 3 x 256 x 256
-    return img / 255.0
-
-def img_preprocessing(img_path, size):
-    img = load_image(img_path, size)  # [0, 1]
-    img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
-    imgs_norm = (img - 0.5) * 2.0  # [-1, 1]
-    return imgs_norm
-
-def saved_image(img_tensor, img_path):
-    toPIL = transforms.ToPILImage()
-    img = toPIL(img_tensor.detach().cpu().squeeze(0))  # 使用squeeze(0)来移除批次维度
-    img.save(img_path)
-
-def main(args):
-    frames_result_saved_path = os.path.join(args.result_path, 'frames')
-    os.makedirs(frames_result_saved_path, exist_ok=True)
-    test_image_name = os.path.splitext(os.path.basename(args.test_image_path))[0]
-    audio_name = os.path.splitext(os.path.basename(args.test_audio_path))[0]
-    predicted_video_256_path = os.path.join(args.result_path,  f'{test_image_name}-{audio_name}.mp4')
-    predicted_video_512_path = os.path.join(args.result_path,  f'{test_image_name}-{audio_name}_SR.mp4')
- 
-    #======Loading Stage 1 model=========
-    lia = LIA_Model(motion_dim=args.motion_dim, fusion_type='weighted_sum')
-    lia.load_lightning_model(args.stage1_checkpoint_path)
-    lia.to(args.device)
-    #============================
-
-    conf = ffhq256_autoenc()
-    conf.seed = args.seed
-    conf.decoder_layers = args.decoder_layers
-    conf.infer_type = args.infer_type
-    conf.motion_dim = args.motion_dim
-    
-    if args.infer_type == 'mfcc_full_control':
-        conf.face_location=True
-        conf.face_scale=True
-        conf.mfcc = True
-    elif args.infer_type == 'mfcc_pose_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = True
-    elif args.infer_type == 'hubert_pose_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = False
-    elif args.infer_type == 'hubert_audio_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = False
-    elif args.infer_type == 'hubert_full_control':
-        conf.face_location=True
-        conf.face_scale=True
-        conf.mfcc = False
-    else:
-        print('Type NOT Found!')
-        exit(0)
-        
-    if not os.path.exists(args.test_image_path):
-        print(f'{args.test_image_path} does not exist!')
-        exit(0)
-    
-    if not os.path.exists(args.test_audio_path):
-        print(f'{args.test_audio_path} does not exist!')
-        exit(0)
-    
-    img_source = img_preprocessing(args.test_image_path, args.image_size).to(args.device)
-    one_shot_lia_start, one_shot_lia_direction, feats = lia.get_start_direction_code(img_source, img_source, img_source, img_source)
-
-    #======Loading Stage 2 model=========
-    model = LitModel(conf)
-    state = torch.load(args.stage2_checkpoint_path, map_location='cpu')
-    model.load_state_dict(state, strict=True)
-    model.ema_model.eval()
-    model.ema_model.to(args.device)
-    #=================================
-    
-    #======Audio Input=========
-    if conf.infer_type.startswith('mfcc'):
-        # MFCC features
-        wav, sr = librosa.load(args.test_audio_path, sr=16000)
-        input_values = python_speech_features.mfcc(signal=wav, samplerate=sr, numcep=13, winlen=0.025, winstep=0.01)
-        d_mfcc_feat = python_speech_features.base.delta(input_values, 1)
-        d_mfcc_feat2 = python_speech_features.base.delta(input_values, 2)
-        audio_driven_obj = np.hstack((input_values, d_mfcc_feat, d_mfcc_feat2))
-        frame_start, frame_end = 0, int(audio_driven_obj.shape[0]/4)
-        audio_start, audio_end = int(frame_start * 4), int(frame_end * 4) # The video frame is fixed to 25 hz and the audio is fixed to 100 hz
-        
-        audio_driven = torch.Tensor(audio_driven_obj[audio_start:audio_end,:]).unsqueeze(0).float().to(args.device)
-        
-    elif conf.infer_type.startswith('hubert'):
-        # Hubert features
-        if not os.path.exists(args.test_hubert_path):
-            
-            if not check_package_installed('transformers'):
-                print('Please install transformers module first.')
-                exit(0)
-            hubert_model_path = './ckpts/chinese-hubert-large'
-            if not os.path.exists(hubert_model_path):
-                print('Please download the hubert weight into the ckpts path first.')
-                exit(0)
-            print('You did not extract the audio features in advance, extracting online now, which will increase processing delay')
-
-            start_time = time.time()
-
-            # load hubert model
-            from transformers import Wav2Vec2FeatureExtractor, HubertModel
-            audio_model = HubertModel.from_pretrained(hubert_model_path).to(args.device)
-            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_path)
-            audio_model.feature_extractor._freeze_parameters()
-            audio_model.eval()
-
-            # hubert model forward pass
-            audio, sr = librosa.load(args.test_audio_path, sr=16000)
-            input_values = feature_extractor(audio, sampling_rate=16000, padding=True, do_normalize=True, return_tensors="pt").input_values
-            input_values = input_values.to(args.device)
-            ws_feats = []
-            with torch.no_grad():
-                outputs = audio_model(input_values, output_hidden_states=True)
-                for i in range(len(outputs.hidden_states)):
-                    ws_feats.append(outputs.hidden_states[i].detach().cpu().numpy())
-                ws_feat_obj = np.array(ws_feats)
-                ws_feat_obj = np.squeeze(ws_feat_obj, 1)
-                ws_feat_obj = np.pad(ws_feat_obj, ((0, 0), (0, 1), (0, 0)), 'edge') # align the audio length with video frame
-            
-            execution_time = time.time() - start_time
-            print(f"Extraction Audio Feature: {execution_time:.2f} Seconds")
-
-            audio_driven_obj = ws_feat_obj
-        else:
-            print(f'Using audio feature from path: {args.test_hubert_path}')
-            audio_driven_obj = np.load(args.test_hubert_path)
-
-        frame_start, frame_end = 0, int(audio_driven_obj.shape[1]/2)
-        audio_start, audio_end = int(frame_start * 2), int(frame_end * 2) # The video frame is fixed to 25 hz and the audio is fixed to 50 hz
-        
-        audio_driven = torch.Tensor(audio_driven_obj[:,audio_start:audio_end,:]).unsqueeze(0).float().to(args.device)
-    #============================
-    
-    # Diffusion Noise
-    noisyT = torch.randn((1,frame_end, args.motion_dim)).to(args.device)
-    
-    #======Inputs for Attribute Control=========
-    if os.path.exists(args.pose_driven_path):
-        pose_obj = np.load(args.pose_driven_path)
-
-        if len(pose_obj.shape) != 2:
-            print('please check your pose information. The shape must be like (T, 3).')
-            exit(0)
-        if pose_obj.shape[1] != 3:
-            print('please check your pose information. The shape must be like (T, 3).')
-            exit(0)
-
-        if pose_obj.shape[0] >= frame_end:
-            pose_obj = pose_obj[:frame_end,:]
-        else:
-            padding = np.tile(pose_obj[-1, :], (frame_end - pose_obj.shape[0], 1))
-            pose_obj = np.vstack((pose_obj, padding))
-            
-        pose_signal = torch.Tensor(pose_obj).unsqueeze(0).to(args.device) / 90 # 90 is for normalization here
-    else:
-        yaw_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_yaw
-        pitch_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_pitch
-        roll_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_roll
-        pose_signal = torch.cat((yaw_signal, pitch_signal, roll_signal), dim=-1)
-    
-    pose_signal = torch.clamp(pose_signal, -1, 1)
-
-    face_location_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.face_location
-    face_scae_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.face_scale
-    #===========================================
-
-    start_time = time.time()
-
-    #======Diffusion Denosing Process=========
-    generated_directions = model.render(one_shot_lia_start, one_shot_lia_direction, audio_driven, face_location_signal, face_scae_signal, pose_signal, noisyT, args.step_T, control_flag=args.control_flag)
-    #=========================================
-    
-    execution_time = time.time() - start_time
-    print(f"Motion Diffusion Model: {execution_time:.2f} Seconds")
-
-    generated_directions = generated_directions.detach().cpu().numpy()
-    
-    start_time = time.time()
-    #======Rendering images frame-by-frame=========
-    for pred_index in tqdm(range(generated_directions.shape[1])):
-        ori_img_recon = lia.render(one_shot_lia_start, torch.Tensor(generated_directions[:,pred_index,:]).to(args.device), feats)
-        ori_img_recon = ori_img_recon.clamp(-1, 1)
-        wav_pred = (ori_img_recon.detach() + 1) / 2
-        saved_image(wav_pred, os.path.join(frames_result_saved_path, "%06d.png"%(pred_index)))
-    #==============================================
-    
-    execution_time = time.time() - start_time
-    print(f"Renderer Model: {execution_time:.2f} Seconds")
-
-    frames_to_video(frames_result_saved_path, args.test_audio_path, predicted_video_256_path)
-    
-    shutil.rmtree(frames_result_saved_path)
-    
-    # Enhancer
-    if args.face_sr and check_package_installed('gfpgan'):
-        from face_sr.face_enhancer import enhancer_list
-        import imageio
-
-        # Super-resolution
-        imageio.mimsave(predicted_video_512_path+'.tmp.mp4', enhancer_list(predicted_video_256_path, method='gfpgan', bg_upsampler=None), fps=float(25))
-        
-        # Merge audio and video
-        video_clip = VideoFileClip(predicted_video_512_path+'.tmp.mp4')
-        audio_clip = AudioFileClip(predicted_video_256_path)
-        final_clip = video_clip.set_audio(audio_clip)
-        final_clip.write_videofile(predicted_video_512_path, codec='libx264', audio_codec='aac')
-        
-        os.remove(predicted_video_512_path+'.tmp.mp4')
-    
-    if args.face_sr:
-        return predicted_video_256_path, predicted_video_512_path
-    else:
-        return predicted_video_256_path, predicted_video_256_path
-
-def generate_video(uploaded_img, uploaded_audio, infer_type, 
-        pose_yaw, pose_pitch, pose_roll, face_location, face_scale, step_T, device, face_sr, seed):
-    if uploaded_img is None or uploaded_audio is None:
-        return None, gr.Markdown("Error: Input image or audio file is empty. Please check and upload both files.")
-    
-    model_mapping = {
-        "mfcc_pose_only": "./ckpts/stage2_pose_only_mfcc.ckpt",
-        "mfcc_full_control": "./ckpts/stage2_more_controllable_mfcc.ckpt",
-        "hubert_audio_only": "./ckpts/stage2_audio_only_hubert.ckpt",
-        "hubert_pose_only": "./ckpts/stage2_pose_only_hubert.ckpt",
-        "hubert_full_control": "./ckpts/stage2_full_control_hubert.ckpt",
-    }
-
-    # if face_crop:
-    #     uploaded_img_path = Path(uploaded_img)
-    #     cropped_img_path = uploaded_img_path.with_name(uploaded_img_path.stem + "_crop" + uploaded_img_path.suffix)
-    #     crop_image(uploaded_img, cropped_img_path)
-    #     uploaded_img = str(cropped_img_path)
-
-    # import pdb;pdb.set_trace()
-
-    stage2_checkpoint_path = model_mapping.get(infer_type, "default_checkpoint.ckpt")
-    try:
-        args = argparse.Namespace(
-            infer_type=infer_type,
-            test_image_path=uploaded_img,
-            test_audio_path=uploaded_audio,
-            test_hubert_path='',
-            result_path='./outputs/',
-            stage1_checkpoint_path='./ckpts/stage1.ckpt',
-            stage2_checkpoint_path=stage2_checkpoint_path,
-            seed=seed,
-            control_flag=True,
-            pose_yaw=pose_yaw,
-            pose_pitch=pose_pitch,
-            pose_roll=pose_roll,
-            face_location=face_location,
-            pose_driven_path='not_supported_in_this_mode',
-            face_scale=face_scale,
-            step_T=step_T,
-            image_size=256,
-            device=device,
-            motion_dim=20,
-            decoder_layers=2,
-            face_sr=face_sr
-        )
-
-        # Save the uploaded audio to the expected path
-        # shutil.copy(uploaded_audio, args.test_audio_path)
-
-        # Run the main function
-        output_256_video_path, output_512_video_path = main(args)
-
-        # Check if the output video file exists
-        if not os.path.exists(output_256_video_path):
-            return None, gr.Markdown("Error: Video generation failed. Please check your inputs and try again.")
-        if output_256_video_path == output_512_video_path:
-            return gr.Video(value=output_256_video_path), None, gr.Markdown("Video (256*256 only) generated successfully!")
-        return gr.Video(value=output_256_video_path), gr.Video(value=output_512_video_path), gr.Markdown("Video generated successfully!")
-
-    except Exception as e:
-        return None, None, gr.Markdown(f"Error: An unexpected error occurred - {str(e)}")
-
-default_values = {
-    "pose_yaw": 0,
-    "pose_pitch": 0,
-    "pose_roll": 0,
-    "face_location": 0.5,
-    "face_scale": 0.5,
-    "step_T": 50,
-    "seed": 0,
-    "device": "cuda"
-}
-
-with gr.Blocks() as demo:
-    gr.Markdown('# AniTalker')
-    gr.Markdown('![]()')
-    with gr.Row():
-        with gr.Column():
-            uploaded_img = gr.Image(type="filepath", label="Reference Image")
-            uploaded_audio = gr.Audio(type="filepath", label="Input Audio")
-        with gr.Column():
-            output_video_256 = gr.Video(label="Generated Video (256)")
-            output_video_512 = gr.Video(label="Generated Video (512)")
-            output_message = gr.Markdown()
-        
-
-
-    generate_button = gr.Button("Generate Video")
-
-    with gr.Accordion("Configuration", open=True):
-        infer_type = gr.Dropdown(
-            label="Inference Type",
-            choices=['mfcc_pose_only', 'mfcc_full_control', 'hubert_audio_only', 'hubert_pose_only'],
-            value='hubert_audio_only'
-        )
-        face_sr = gr.Checkbox(label="Enable Face Super-Resolution (512*512)", value=False)
-        # face_crop = gr.Checkbox(label="Face Crop (Dlib)", value=False) 
-        # face_crop = False # TODO
-        seed = gr.Number(label="Seed", value=default_values["seed"])
-        pose_yaw = gr.Slider(label="pose_yaw", minimum=-1, maximum=1, value=default_values["pose_yaw"])
-        pose_pitch = gr.Slider(label="pose_pitch", minimum=-1, maximum=1, value=default_values["pose_pitch"])
-        pose_roll = gr.Slider(label="pose_roll", minimum=-1, maximum=1, value=default_values["pose_roll"])
-        face_location = gr.Slider(label="face_location", minimum=0, maximum=1, value=default_values["face_location"])
-        face_scale = gr.Slider(label="face_scale", minimum=0, maximum=1, value=default_values["face_scale"])
-        step_T = gr.Slider(label="step_T", minimum=1, maximum=100, step=1, value=default_values["step_T"])
-        device = gr.Radio(label="Device", choices=["cuda", "cpu"], value=default_values["device"])
-        
-
-    generate_button.click(
-        generate_video,
-        inputs=[
-            uploaded_img, uploaded_audio, infer_type, 
-            pose_yaw, pose_pitch, pose_roll, face_location, face_scale, step_T, device, face_sr, seed
-        ],
-        outputs=[output_video_256, output_video_512, output_message]
-    )
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='EchoMimic')
-    parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Server name')
-    parser.add_argument('--server_port', type=int, default=3001, help='Server port')
-    args = parser.parse_args()
-
-    demo.launch()

code/choices.py

@@ -1,179 +0,0 @@
-from enum import Enum
-from torch import nn
-
-
-class TrainMode(Enum):
-    # manipulate mode = training the classifier
-    manipulate = 'manipulate'
-    # default trainin mode!
-    diffusion = 'diffusion'
-    # default latent training mode!
-    # fitting the a DDPM to a given latent
-    latent_diffusion = 'latentdiffusion'
-
-    def is_manipulate(self):
-        return self in [
-            TrainMode.manipulate,
-        ]
-
-    def is_diffusion(self):
-        return self in [
-            TrainMode.diffusion,
-            TrainMode.latent_diffusion,
-        ]
-
-    def is_autoenc(self):
-        # the network possibly does autoencoding
-        return self in [
-            TrainMode.diffusion,
-        ]
-
-    def is_latent_diffusion(self):
-        return self in [
-            TrainMode.latent_diffusion,
-        ]
-
-    def use_latent_net(self):
-        return self.is_latent_diffusion()
-
-    def require_dataset_infer(self):
-        """
-        whether training in this mode requires the latent variables to be available?
-        """
-        # this will precalculate all the latents before hand
-        # and the dataset will be all the predicted latents
-        return self in [
-            TrainMode.latent_diffusion,
-            TrainMode.manipulate,
-        ]
-
-
-class ManipulateMode(Enum):
-    """
-    how to train the classifier to manipulate
-    """
-    # train on whole celeba attr dataset
-    celebahq_all = 'celebahq_all'
-    # celeba with D2C's crop
-    d2c_fewshot = 'd2cfewshot'
-    d2c_fewshot_allneg = 'd2cfewshotallneg'
-
-    def is_celeba_attr(self):
-        return self in [
-            ManipulateMode.d2c_fewshot,
-            ManipulateMode.d2c_fewshot_allneg,
-            ManipulateMode.celebahq_all,
-        ]
-
-    def is_single_class(self):
-        return self in [
-            ManipulateMode.d2c_fewshot,
-            ManipulateMode.d2c_fewshot_allneg,
-        ]
-
-    def is_fewshot(self):
-        return self in [
-            ManipulateMode.d2c_fewshot,
-            ManipulateMode.d2c_fewshot_allneg,
-        ]
-
-    def is_fewshot_allneg(self):
-        return self in [
-            ManipulateMode.d2c_fewshot_allneg,
-        ]
-
-
-class ModelType(Enum):
-    """
-    Kinds of the backbone models
-    """
-
-    # unconditional ddpm
-    ddpm = 'ddpm'
-    # autoencoding ddpm cannot do unconditional generation
-    autoencoder = 'autoencoder'
-
-    def has_autoenc(self):
-        return self in [
-            ModelType.autoencoder,
-        ]
-
-    def can_sample(self):
-        return self in [ModelType.ddpm]
-
-
-class ModelName(Enum):
-    """
-    List of all supported model classes
-    """
-
-    beatgans_ddpm = 'beatgans_ddpm'
-    beatgans_autoenc = 'beatgans_autoenc'
-
-
-class ModelMeanType(Enum):
-    """
-    Which type of output the model predicts.
-    """
-
-    eps = 'eps'  # the model predicts epsilon
-
-
-class ModelVarType(Enum):
-    """
-    What is used as the model's output variance.
-
-    The LEARNED_RANGE option has been added to allow the model to predict
-    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
-    """
-
-    # posterior beta_t
-    fixed_small = 'fixed_small'
-    # beta_t
-    fixed_large = 'fixed_large'
-
-
-class LossType(Enum):
-    mse = 'mse'  # use raw MSE loss (and KL when learning variances)
-    l1 = 'l1'
-
-
-class GenerativeType(Enum):
-    """
-    How's a sample generated
-    """
-
-    ddpm = 'ddpm'
-    ddim = 'ddim'
-
-
-class OptimizerType(Enum):
-    adam = 'adam'
-    adamw = 'adamw'
-
-
-class Activation(Enum):
-    none = 'none'
-    relu = 'relu'
-    lrelu = 'lrelu'
-    silu = 'silu'
-    tanh = 'tanh'
-
-    def get_act(self):
-        if self == Activation.none:
-            return nn.Identity()
-        elif self == Activation.relu:
-            return nn.ReLU()
-        elif self == Activation.lrelu:
-            return nn.LeakyReLU(negative_slope=0.2)
-        elif self == Activation.silu:
-            return nn.SiLU()
-        elif self == Activation.tanh:
-            return nn.Tanh()
-        else:
-            raise NotImplementedError()
-
-
-class ManipulateLossType(Enum):
-    bce = 'bce'
-    mse = 'mse'

code/config.py

@@ -1,388 +0,0 @@
-from model.unet import ScaleAt
-from model.latentnet import *
-from diffusion.resample import UniformSampler
-from diffusion.diffusion import space_timesteps
-from typing import Tuple
-
-from torch.utils.data import DataLoader
-
-from config_base import BaseConfig
-from diffusion import *
-from diffusion.base import GenerativeType, LossType, ModelMeanType, ModelVarType, get_named_beta_schedule
-from model import *
-from choices import *
-from multiprocessing import get_context
-import os
-from dataset_util import *
-from torch.utils.data.distributed import DistributedSampler
-from dataset import LatentDataLoader
-
-@dataclass
-class PretrainConfig(BaseConfig):
-    name: str
-    path: str
-
-
-@dataclass
-class TrainConfig(BaseConfig):
-    # random seed
-    seed: int = 0
-    train_mode: TrainMode = TrainMode.diffusion
-    train_cond0_prob: float = 0
-    train_pred_xstart_detach: bool = True
-    train_interpolate_prob: float = 0
-    train_interpolate_img: bool = False
-    manipulate_mode: ManipulateMode = ManipulateMode.celebahq_all
-    manipulate_cls: str = None
-    manipulate_shots: int = None
-    manipulate_loss: ManipulateLossType = ManipulateLossType.bce
-    manipulate_znormalize: bool = False
-    manipulate_seed: int = 0
-    accum_batches: int = 1
-    autoenc_mid_attn: bool = True
-    batch_size: int = 16
-    batch_size_eval: int = None
-    beatgans_gen_type: GenerativeType = GenerativeType.ddim
-    beatgans_loss_type: LossType = LossType.mse
-    beatgans_model_mean_type: ModelMeanType = ModelMeanType.eps
-    beatgans_model_var_type: ModelVarType = ModelVarType.fixed_large
-    beatgans_rescale_timesteps: bool = False
-    latent_infer_path: str = None
-    latent_znormalize: bool = False
-    latent_gen_type: GenerativeType = GenerativeType.ddim
-    latent_loss_type: LossType = LossType.mse
-    latent_model_mean_type: ModelMeanType = ModelMeanType.eps
-    latent_model_var_type: ModelVarType = ModelVarType.fixed_large
-    latent_rescale_timesteps: bool = False
-    latent_T_eval: int = 1_000
-    latent_clip_sample: bool = False
-    latent_beta_scheduler: str = 'linear'
-    beta_scheduler: str = 'linear'
-    data_name: str = ''
-    data_val_name: str = None
-    diffusion_type: str = None
-    dropout: float = 0.1
-    ema_decay: float = 0.9999
-    eval_num_images: int = 5_000
-    eval_every_samples: int = 200_000
-    eval_ema_every_samples: int = 200_000
-    fid_use_torch: bool = True
-    fp16: bool = False
-    grad_clip: float = 1
-    img_size: int = 64
-    lr: float = 0.0001
-    optimizer: OptimizerType = OptimizerType.adam
-    weight_decay: float = 0
-    model_conf: ModelConfig = None
-    model_name: ModelName = None
-    model_type: ModelType = None
-    net_attn: Tuple[int] = None
-    net_beatgans_attn_head: int = 1
-    # not necessarily the same as the the number of style channels
-    net_beatgans_embed_channels: int = 512
-    net_resblock_updown: bool = True
-    net_enc_use_time: bool = False
-    net_enc_pool: str = 'adaptivenonzero'
-    net_beatgans_gradient_checkpoint: bool = False
-    net_beatgans_resnet_two_cond: bool = False
-    net_beatgans_resnet_use_zero_module: bool = True
-    net_beatgans_resnet_scale_at: ScaleAt = ScaleAt.after_norm
-    net_beatgans_resnet_cond_channels: int = None
-    net_ch_mult: Tuple[int] = None
-    net_ch: int = 64
-    net_enc_attn: Tuple[int] = None
-    net_enc_k: int = None
-    # number of resblocks for the encoder (half-unet)
-    net_enc_num_res_blocks: int = 2
-    net_enc_channel_mult: Tuple[int] = None
-    net_enc_grad_checkpoint: bool = False
-    net_autoenc_stochastic: bool = False
-    net_latent_activation: Activation = Activation.silu
-    net_latent_channel_mult: Tuple[int] = (1, 2, 4)
-    net_latent_condition_bias: float = 0
-    net_latent_dropout: float = 0
-    net_latent_layers: int = None
-    net_latent_net_last_act: Activation = Activation.none
-    net_latent_net_type: LatentNetType = LatentNetType.none
-    net_latent_num_hid_channels: int = 1024
-    net_latent_num_time_layers: int = 2
-    net_latent_skip_layers: Tuple[int] = None
-    net_latent_time_emb_channels: int = 64
-    net_latent_use_norm: bool = False
-    net_latent_time_last_act: bool = False
-    net_num_res_blocks: int = 2
-    # number of resblocks for the UNET
-    net_num_input_res_blocks: int = None
-    net_enc_num_cls: int = None
-    num_workers: int = 4
-    parallel: bool = False
-    postfix: str = ''
-    sample_size: int = 64
-    sample_every_samples: int = 20_000
-    save_every_samples: int = 100_000
-    style_ch: int = 512
-    T_eval: int = 1_000
-    T_sampler: str = 'uniform'
-    T: int = 1_000
-    total_samples: int = 10_000_000
-    warmup: int = 0
-    pretrain: PretrainConfig = None
-    continue_from: PretrainConfig = None
-    eval_programs: Tuple[str] = None
-    # if present load the checkpoint from this path instead
-    eval_path: str = None
-    base_dir: str = 'checkpoints'
-    use_cache_dataset: bool = False
-    data_cache_dir: str = os.path.expanduser('~/cache')
-    work_cache_dir: str = os.path.expanduser('~/mycache')
-    # to be overridden
-    name: str = ''
-
-    def __post_init__(self):
-        self.batch_size_eval = self.batch_size_eval or self.batch_size
-        self.data_val_name = self.data_val_name or self.data_name
-
-    def scale_up_gpus(self, num_gpus, num_nodes=1):
-        self.eval_ema_every_samples *= num_gpus * num_nodes
-        self.eval_every_samples *= num_gpus * num_nodes
-        self.sample_every_samples *= num_gpus * num_nodes
-        self.batch_size *= num_gpus * num_nodes
-        self.batch_size_eval *= num_gpus * num_nodes
-        return self
-
-    @property
-    def batch_size_effective(self):
-        return self.batch_size * self.accum_batches
-
-    @property
-    def fid_cache(self):
-        # we try to use the local dirs to reduce the load over network drives
-        # hopefully, this would reduce the disconnection problems with sshfs
-        return f'{self.work_cache_dir}/eval_images/{self.data_name}_size{self.img_size}_{self.eval_num_images}'
-
-    @property
-    def data_path(self):
-        # may use the cache dir
-        path = data_paths[self.data_name]
-        if self.use_cache_dataset and path is not None:
-            path = use_cached_dataset_path(
-                path, f'{self.data_cache_dir}/{self.data_name}')
-        return path
-
-    @property
-    def logdir(self):
-        return f'{self.base_dir}/{self.name}'
-
-    @property
-    def generate_dir(self):
-        # we try to use the local dirs to reduce the load over network drives
-        # hopefully, this would reduce the disconnection problems with sshfs
-        return f'{self.work_cache_dir}/gen_images/{self.name}'
-
-    def _make_diffusion_conf(self, T=None):
-        if self.diffusion_type == 'beatgans':
-            # can use T < self.T for evaluation
-            # follows the guided-diffusion repo conventions
-            # t's are evenly spaced
-            if self.beatgans_gen_type == GenerativeType.ddpm:
-                section_counts = [T]
-            elif self.beatgans_gen_type == GenerativeType.ddim:
-                section_counts = f'ddim{T}'
-            else:
-                raise NotImplementedError()
-
-            return SpacedDiffusionBeatGansConfig(
-                gen_type=self.beatgans_gen_type,
-                model_type=self.model_type,
-                betas=get_named_beta_schedule(self.beta_scheduler, self.T),
-                model_mean_type=self.beatgans_model_mean_type,
-                model_var_type=self.beatgans_model_var_type,
-                loss_type=self.beatgans_loss_type,
-                rescale_timesteps=self.beatgans_rescale_timesteps,
-                use_timesteps=space_timesteps(num_timesteps=self.T,
-                                              section_counts=section_counts),
-                fp16=self.fp16,
-            )
-        else:
-            raise NotImplementedError()
-
-    def _make_latent_diffusion_conf(self, T=None):
-        # can use T < self.T for evaluation
-        # follows the guided-diffusion repo conventions
-        # t's are evenly spaced
-        if self.latent_gen_type == GenerativeType.ddpm:
-            section_counts = [T]
-        elif self.latent_gen_type == GenerativeType.ddim:
-            section_counts = f'ddim{T}'
-        else:
-            raise NotImplementedError()
-
-        return SpacedDiffusionBeatGansConfig(
-            train_pred_xstart_detach=self.train_pred_xstart_detach,
-            gen_type=self.latent_gen_type,
-            # latent's model is always ddpm
-            model_type=ModelType.ddpm,
-            # latent shares the beta scheduler and full T
-            betas=get_named_beta_schedule(self.latent_beta_scheduler, self.T),
-            model_mean_type=self.latent_model_mean_type,
-            model_var_type=self.latent_model_var_type,
-            loss_type=self.latent_loss_type,
-            rescale_timesteps=self.latent_rescale_timesteps,
-            use_timesteps=space_timesteps(num_timesteps=self.T,
-                                          section_counts=section_counts),
-            fp16=self.fp16,
-        )
-
-    @property
-    def model_out_channels(self):
-        return 3
-
-    def make_T_sampler(self):
-        if self.T_sampler == 'uniform':
-            return UniformSampler(self.T)
-        else:
-            raise NotImplementedError()
-
-    def make_diffusion_conf(self):
-        return self._make_diffusion_conf(self.T)
-
-    def make_eval_diffusion_conf(self):
-        return self._make_diffusion_conf(T=self.T_eval)
-
-    def make_latent_diffusion_conf(self):
-        return self._make_latent_diffusion_conf(T=self.T)
-
-    def make_latent_eval_diffusion_conf(self):
-        # latent can have different eval T
-        return self._make_latent_diffusion_conf(T=self.latent_T_eval)
-
-    def make_dataset(self, path=None, **kwargs):
-        return LatentDataLoader(self.window_size, 
-                                self.frame_jpgs,
-                                self.lmd_feats_prefix, 
-                                self.audio_prefix, 
-                                self.raw_audio_prefix, 
-                                self.motion_latents_prefix, 
-                                self.pose_prefix,
-                                self.db_name,
-                                audio_hz=self.audio_hz)
-
-    def make_loader(self,
-                    dataset,
-                    shuffle: bool,
-                    num_worker: bool = None,
-                    drop_last: bool = True,
-                    batch_size: int = None,
-                    parallel: bool = False):
-        if parallel and distributed.is_initialized():
-            # drop last to make sure that there is no added special indexes
-            sampler = DistributedSampler(dataset,
-                                         shuffle=shuffle,
-                                         drop_last=True)
-        else:
-            sampler = None
-        return DataLoader(
-            dataset,
-            batch_size=batch_size or self.batch_size,
-            sampler=sampler,
-            # with sampler, use the sample instead of this option
-            shuffle=False if sampler else shuffle,
-            num_workers=num_worker or self.num_workers,
-            pin_memory=True,
-            drop_last=drop_last,
-            multiprocessing_context=get_context('fork'),
-        )
-
-    def make_model_conf(self):
-        if self.model_name == ModelName.beatgans_ddpm:
-            self.model_type = ModelType.ddpm
-            self.model_conf = BeatGANsUNetConfig(
-                attention_resolutions=self.net_attn,
-                channel_mult=self.net_ch_mult,
-                conv_resample=True,
-                dims=2,
-                dropout=self.dropout,
-                embed_channels=self.net_beatgans_embed_channels,
-                image_size=self.img_size,
-                in_channels=3,
-                model_channels=self.net_ch,
-                num_classes=None,
-                num_head_channels=-1,
-                num_heads_upsample=-1,
-                num_heads=self.net_beatgans_attn_head,
-                num_res_blocks=self.net_num_res_blocks,
-                num_input_res_blocks=self.net_num_input_res_blocks,
-                out_channels=self.model_out_channels,
-                resblock_updown=self.net_resblock_updown,
-                use_checkpoint=self.net_beatgans_gradient_checkpoint,
-                use_new_attention_order=False,
-                resnet_two_cond=self.net_beatgans_resnet_two_cond,
-                resnet_use_zero_module=self.
-                net_beatgans_resnet_use_zero_module,
-            )
-        elif self.model_name in [
-                ModelName.beatgans_autoenc,
-        ]:
-            cls = BeatGANsAutoencConfig
-            # supports both autoenc and vaeddpm
-            if self.model_name == ModelName.beatgans_autoenc:
-                self.model_type = ModelType.autoencoder
-            else:
-                raise NotImplementedError()
-
-            if self.net_latent_net_type == LatentNetType.none:
-                latent_net_conf = None
-            elif self.net_latent_net_type == LatentNetType.skip:
-                latent_net_conf = MLPSkipNetConfig(
-                    num_channels=self.style_ch,
-                    skip_layers=self.net_latent_skip_layers,
-                    num_hid_channels=self.net_latent_num_hid_channels,
-                    num_layers=self.net_latent_layers,
-                    num_time_emb_channels=self.net_latent_time_emb_channels,
-                    activation=self.net_latent_activation,
-                    use_norm=self.net_latent_use_norm,
-                    condition_bias=self.net_latent_condition_bias,
-                    dropout=self.net_latent_dropout,
-                    last_act=self.net_latent_net_last_act,
-                    num_time_layers=self.net_latent_num_time_layers,
-                    time_last_act=self.net_latent_time_last_act,
-                )
-            else:
-                raise NotImplementedError()
-
-            self.model_conf = cls(
-                attention_resolutions=self.net_attn,
-                channel_mult=self.net_ch_mult,
-                conv_resample=True,
-                dims=2,
-                dropout=self.dropout,
-                embed_channels=self.net_beatgans_embed_channels,
-                enc_out_channels=self.style_ch,
-                enc_pool=self.net_enc_pool,
-                enc_num_res_block=self.net_enc_num_res_blocks,
-                enc_channel_mult=self.net_enc_channel_mult,
-                enc_grad_checkpoint=self.net_enc_grad_checkpoint,
-                enc_attn_resolutions=self.net_enc_attn,
-                image_size=self.img_size,
-                in_channels=3,
-                model_channels=self.net_ch,
-                num_classes=None,
-                num_head_channels=-1,
-                num_heads_upsample=-1,
-                num_heads=self.net_beatgans_attn_head,
-                num_res_blocks=self.net_num_res_blocks,
-                num_input_res_blocks=self.net_num_input_res_blocks,
-                out_channels=self.model_out_channels,
-                resblock_updown=self.net_resblock_updown,
-                use_checkpoint=self.net_beatgans_gradient_checkpoint,
-                use_new_attention_order=False,
-                resnet_two_cond=self.net_beatgans_resnet_two_cond,
-                resnet_use_zero_module=self.
-                net_beatgans_resnet_use_zero_module,
-                latent_net_conf=latent_net_conf,
-                resnet_cond_channels=self.net_beatgans_resnet_cond_channels,
-            )
-        else:
-            raise NotImplementedError(self.model_name)
-
-        return self.model_conf

code/config_base.py

@@ -1,72 +0,0 @@
-import json
-import os
-from copy import deepcopy
-from dataclasses import dataclass
-
-
-@dataclass
-class BaseConfig:
-    def clone(self):
-        return deepcopy(self)
-
-    def inherit(self, another):
-        """inherit common keys from a given config"""
-        common_keys = set(self.__dict__.keys()) & set(another.__dict__.keys())
-        for k in common_keys:
-            setattr(self, k, getattr(another, k))
-
-    def propagate(self):
-        """push down the configuration to all members"""
-        for k, v in self.__dict__.items():
-            if isinstance(v, BaseConfig):
-                v.inherit(self)
-                v.propagate()
-
-    def save(self, save_path):
-        """save config to json file"""
-        dirname = os.path.dirname(save_path)
-        if not os.path.exists(dirname):
-            os.makedirs(dirname)
-        conf = self.as_dict_jsonable()
-        with open(save_path, 'w') as f:
-            json.dump(conf, f)
-
-    def load(self, load_path):
-        """load json config"""
-        with open(load_path) as f:
-            conf = json.load(f)
-        self.from_dict(conf)
-
-    def from_dict(self, dict, strict=False):
-        for k, v in dict.items():
-            if not hasattr(self, k):
-                if strict:
-                    raise ValueError(f"loading extra '{k}'")
-                else:
-                    print(f"loading extra '{k}'")
-                    continue
-            if isinstance(self.__dict__[k], BaseConfig):
-                self.__dict__[k].from_dict(v)
-            else:
-                self.__dict__[k] = v
-
-    def as_dict_jsonable(self):
-        conf = {}
-        for k, v in self.__dict__.items():
-            if isinstance(v, BaseConfig):
-                conf[k] = v.as_dict_jsonable()
-            else:
-                if jsonable(v):
-                    conf[k] = v
-                else:
-                    # ignore not jsonable
-                    pass
-        return conf
-
-
-def jsonable(x):
-    try:
-        json.dumps(x)
-        return True
-    except TypeError:
-        return False

code/dataset.py

@@ -1,218 +0,0 @@
-import os
-import librosa
-from PIL import Image
-from torchvision import transforms
-import python_speech_features
-import random
-import os
-import numpy as np
-from tqdm import tqdm
-import torchvision
-import torchvision.transforms as transforms
-from PIL import Image
-
-class LatentDataLoader(object):
-
-    def __init__(
-        self, 
-        window_size, 
-        frame_jpgs, 
-        lmd_feats_prefix, 
-        audio_prefix, 
-        raw_audio_prefix,
-        motion_latents_prefix, 
-        pose_prefix, 
-        db_name, 
-        video_fps=25, 
-        audio_hz=50, 
-        size=256, 
-        mfcc_mode=False,
-    ):
-        self.window_size = window_size
-        self.lmd_feats_prefix = lmd_feats_prefix
-        self.audio_prefix = audio_prefix
-        self.pose_prefix = pose_prefix
-        self.video_fps = video_fps
-        self.audio_hz = audio_hz
-        self.db_name = db_name
-        self.raw_audio_prefix = raw_audio_prefix
-        self.mfcc_mode = mfcc_mode
-        
-
-        self.transform = torchvision.transforms.Compose([
-            transforms.Resize((size, size)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))]
-        )
-    
-        self.data = []
-        for db_name in [ 'VoxCeleb2', 'HDTF' ]:
-            db_png_path = os.path.join(frame_jpgs, db_name)
-            for clip_name in tqdm(os.listdir(db_png_path)):
-
-                item_dict = dict()
-                item_dict['clip_name'] = clip_name
-                item_dict['frame_count'] =  len(list(os.listdir(os.path.join(frame_jpgs, db_name, clip_name))))
-                item_dict['hubert_path'] = os.path.join(audio_prefix, db_name, clip_name +".npy")
-                item_dict['wav_path'] = os.path.join(raw_audio_prefix, db_name, clip_name +".wav")
-                
-                item_dict['yaw_pitch_roll_path'] = os.path.join(pose_prefix, db_name, 'raw_videos_pose_yaw_pitch_roll', clip_name +".npy")
-                if not os.path.exists(item_dict['yaw_pitch_roll_path']):
-                    print(f"{db_name}'s {clip_name} miss yaw_pitch_roll_path")
-                    continue
-
-                item_dict['yaw_pitch_roll'] = np.load(item_dict['yaw_pitch_roll_path'])
-                item_dict['yaw_pitch_roll'] = np.clip(item_dict['yaw_pitch_roll'], -90, 90) / 90.0
-                
-                if not os.path.exists(item_dict['wav_path']):
-                    print(f"{db_name}'s {clip_name} miss wav_path")
-                    continue 
-                
-                if not os.path.exists(item_dict['hubert_path']):
-                    print(f"{db_name}'s {clip_name} miss hubert_path")
-                    continue 
-                
-                
-                if self.mfcc_mode:
-                    wav, sr = librosa.load(item_dict['wav_path'], sr=16000)
-                    input_values = python_speech_features.mfcc(signal=wav,samplerate=sr,numcep=13,winlen=0.025,winstep=0.01)
-                    d_mfcc_feat = python_speech_features.base.delta(input_values, 1)
-                    d_mfcc_feat2 = python_speech_features.base.delta(input_values, 2)
-                    input_values = np.hstack((input_values, d_mfcc_feat, d_mfcc_feat2))
-                    item_dict['hubert_obj'] = input_values
-                else:
-                    item_dict['hubert_obj'] = np.load(item_dict['hubert_path'], mmap_mode='r')
-                item_dict['lmd_path'] = os.path.join(lmd_feats_prefix, db_name, clip_name +".txt")
-                item_dict['lmd_obj_full'] = self.read_landmark_info(item_dict['lmd_path'], upper_face=False)
-
-                motion_start_path = os.path.join(motion_latents_prefix, db_name, 'motions', clip_name +".npy")
-                motion_direction_path = os.path.join(motion_latents_prefix, db_name, 'directions',  clip_name +".npy")
-                
-                if not os.path.exists(motion_start_path):
-                    print(f"{db_name}'s {clip_name} miss motion_start_path")
-                    continue
-                if not os.path.exists(motion_direction_path):
-                    print(f"{db_name}'s {clip_name} miss motion_direction_path")
-                    continue
-
-                item_dict['motion_start_obj'] = np.load(motion_start_path)
-                item_dict['motion_direction_obj'] = np.load(motion_direction_path)                
-
-                if self.mfcc_mode:
-                    min_len = min(
-                        item_dict['lmd_obj_full'].shape[0],
-                        item_dict['yaw_pitch_roll'].shape[0],
-                        item_dict['motion_start_obj'].shape[0],
-                        item_dict['motion_direction_obj'].shape[0],
-                        int(item_dict['hubert_obj'].shape[0]/4),
-                        item_dict['frame_count']
-                    )
-                    item_dict['frame_count'] = min_len
-                    item_dict['hubert_obj'] =  item_dict['hubert_obj'][:min_len*4,:]
-                else:
-                    min_len = min(
-                        item_dict['lmd_obj_full'].shape[0],
-                        item_dict['yaw_pitch_roll'].shape[0],
-                        item_dict['motion_start_obj'].shape[0],
-                        item_dict['motion_direction_obj'].shape[0],
-                        int(item_dict['hubert_obj'].shape[1]/2),
-                        item_dict['frame_count']
-                    )
-                
-                    item_dict['frame_count'] = min_len
-                    item_dict['hubert_obj'] =  item_dict['hubert_obj'][:, :min_len*2, :]
-                
-                if min_len < self.window_size * self.video_fps + 5:
-                    continue
-
-        print('Db count:', len(self.data))
-
-    def get_single_image(self, image_path):
-        img_source = Image.open(image_path).convert('RGB')
-        img_source = self.transform(img_source)
-        return img_source
-
-    def get_multiple_ranges(self, lists, multi_ranges):
-        # Ensure that multi_ranges is a list of tuples
-        if not all(isinstance(item, tuple) and len(item) == 2 for item in multi_ranges):
-            raise ValueError("multi_ranges must be a list of (start, end) tuples with exactly two elements each")
-        extracted_elements = [lists[start:end] for start, end in multi_ranges]
-        flat_list = [item for sublist in extracted_elements for item in sublist]
-        return flat_list
-    
-    
-    def read_landmark_info(self, lmd_path, upper_face=True):
-        with open(lmd_path, 'r') as file:
-            lmd_lines = file.readlines()
-        lmd_lines.sort()
-        
-        total_lmd_obj = []
-        for i, line in enumerate(lmd_lines):
-            # Split the coordinates and filter out any empty strings
-            coords = [c for c in line.strip().split(' ') if c]
-            coords = coords[1:] # do not include the file name in the first row
-            lmd_obj = []
-            if upper_face:
-                # Ensure that the coordinates are parsed as integers
-                for coord_pair in self.get_multiple_ranges(coords, [(0, 3), (14, 27), (36, 48)]): # 28个
-                    x, y = coord_pair.split('_')
-                    lmd_obj.append((int(x)/512, int(y)/512))
-            else:
-                for coord_pair in coords:
-                    x, y = coord_pair.split('_')
-                    lmd_obj.append((int(x)/512, int(y)/512))
-            total_lmd_obj.append(lmd_obj)
-        
-        return np.array(total_lmd_obj, dtype=np.float32)
-
-    def calculate_face_height(self, landmarks):
-        forehead_center = (landmarks[ :, 21, :] + landmarks[:, 22, :]) / 2
-        chin_bottom = landmarks[:, 8, :]
-        distances = np.linalg.norm(forehead_center - chin_bottom, axis=1, keepdims=True)
-        return distances
-            
-    def __getitem__(self, index):
-        
-        data_item = self.data[index]
-        hubert_obj = data_item['hubert_obj']
-        frame_count = data_item['frame_count']
-        lmd_obj_full = data_item['lmd_obj_full']
-        yaw_pitch_roll = data_item['yaw_pitch_roll']
-        motion_start_obj = data_item['motion_start_obj']
-        motion_direction_obj = data_item['motion_direction_obj']
- 
-        frame_end_index = random.randint(self.window_size * self.video_fps + 1, frame_count - 1)
-        frame_start_index = frame_end_index - self.window_size * self.video_fps
-        frame_hint_index = frame_start_index - 1
-        
-        audio_start_index = int(frame_start_index * (self.audio_hz / self.video_fps))
-        audio_end_index = int(frame_end_index * (self.audio_hz / self.video_fps))
-        
-        if self.mfcc_mode:
-            audio_feats = hubert_obj[audio_start_index:audio_end_index, :]
-        else:
-            audio_feats = hubert_obj[:, audio_start_index:audio_end_index, :]
-
-        lmd_obj_full = lmd_obj_full[frame_hint_index:frame_end_index, :]
-        
-        yaw_pitch_roll = yaw_pitch_roll[frame_start_index:frame_end_index, :]
-
-        motion_start = motion_start_obj[frame_hint_index]
-        motion_direction_start = motion_direction_obj[frame_hint_index]
-        motion_direction = motion_direction_obj[frame_start_index:frame_end_index, :]
-
-
-
-        return {
-            'motion_start': motion_start,
-            'motion_direction': motion_direction,
-            'audio_feats': audio_feats,
-            'face_location': lmd_obj_full[1:, 30, 0], # '1:' means taking the first frame as the driven frame. '30' is the noise location, '0' means x coordinate
-            'face_scale': self.calculate_face_height(lmd_obj_full[1:,:,:]),
-            'yaw_pitch_roll': yaw_pitch_roll,
-            'motion_direction_start': motion_direction_start,
-        }
-  
-    def __len__(self):
-        return len(self.data)
-    

code/dataset_util.py

@@ -1,13 +0,0 @@
-import shutil
-import os
-from dist_utils import *
-
-
-def use_cached_dataset_path(source_path, cache_path):
-    if get_rank() == 0:
-        if not os.path.exists(cache_path):
-            # shutil.rmtree(cache_path)
-            print(f'copying the data: {source_path} to {cache_path}')
-            shutil.copytree(source_path, cache_path)
-    barrier()
-    return cache_path

code/demo.py

@@ -1,295 +0,0 @@
-from LIA_Model import LIA_Model
-import torch
-import numpy as np
-import os
-from PIL import Image
-from tqdm import tqdm
-import argparse
-import numpy as np
-from torchvision import transforms
-from templates import *
-import argparse
-import shutil
-from moviepy.editor import *
-import librosa
-import python_speech_features
-import importlib.util
-import time
-
-def check_package_installed(package_name):
-    package_spec = importlib.util.find_spec(package_name)
-    if package_spec is None:
-        print(f"{package_name} is not installed.")
-        return False
-    else:
-        print(f"{package_name} is installed.")
-        return True
-
-def frames_to_video(input_path, audio_path, output_path, fps=25):
-    image_files = [os.path.join(input_path, img) for img in sorted(os.listdir(input_path))]
-    clips = [ImageClip(m).set_duration(1/fps) for m in image_files]
-    video = concatenate_videoclips(clips, method="compose")
-
-    audio = AudioFileClip(audio_path)
-    final_video = video.set_audio(audio)
-    final_video.write_videofile(output_path, fps=fps, codec='libx264', audio_codec='aac')
-
-def load_image(filename, size):
-    img = Image.open(filename).convert('RGB')
-    img = img.resize((size, size))
-    img = np.asarray(img)
-    img = np.transpose(img, (2, 0, 1))  # 3 x 256 x 256
-    return img / 255.0
-
-def img_preprocessing(img_path, size):
-    img = load_image(img_path, size)  # [0, 1]
-    img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
-    imgs_norm = (img - 0.5) * 2.0  # [-1, 1]
-    return imgs_norm
-
-def saved_image(img_tensor, img_path):
-    toPIL = transforms.ToPILImage()
-    img = toPIL(img_tensor.detach().cpu().squeeze(0))  # 使用squeeze(0)来移除批次维度
-    img.save(img_path)
-
-def main(args):
-    frames_result_saved_path = os.path.join(args.result_path, 'frames')
-    os.makedirs(frames_result_saved_path, exist_ok=True)
-    test_image_name = os.path.splitext(os.path.basename(args.test_image_path))[0]
-    audio_name = os.path.splitext(os.path.basename(args.test_audio_path))[0]
-    predicted_video_256_path = os.path.join(args.result_path,  f'{test_image_name}-{audio_name}.mp4')
-    predicted_video_512_path = os.path.join(args.result_path,  f'{test_image_name}-{audio_name}_SR.mp4')
- 
-    #======Loading Stage 1 model=========
-    lia = LIA_Model(motion_dim=args.motion_dim, fusion_type='weighted_sum')
-    lia.load_lightning_model(args.stage1_checkpoint_path)
-    lia.to(args.device)
-    #============================
-
-    conf = ffhq256_autoenc()
-    conf.seed = args.seed
-    conf.decoder_layers = args.decoder_layers
-    conf.infer_type = args.infer_type
-    conf.motion_dim = args.motion_dim
-    
-    if args.infer_type == 'mfcc_full_control':
-        conf.face_location=True
-        conf.face_scale=True
-        conf.mfcc = True
-
-    elif args.infer_type == 'mfcc_pose_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = True
-
-    elif args.infer_type == 'hubert_pose_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = False
-
-    elif args.infer_type == 'hubert_audio_only':
-        conf.face_location=False
-        conf.face_scale=False
-        conf.mfcc = False
-
-    elif args.infer_type == 'hubert_full_control':
-        conf.face_location=True
-        conf.face_scale=True
-        conf.mfcc = False
-
-    else:
-        print('Type NOT Found!')
-        exit(0)
-        
-    if not os.path.exists(args.test_image_path):
-        print(f'{args.test_image_path} does not exist!')
-        exit(0)
-    
-    if not os.path.exists(args.test_audio_path):
-        print(f'{args.test_audio_path} does not exist!')
-        exit(0)
-    
-    img_source = img_preprocessing(args.test_image_path, args.image_size).to(args.device)
-    one_shot_lia_start, one_shot_lia_direction, feats = lia.get_start_direction_code(img_source, img_source, img_source, img_source)
-
-
-    #======Loading Stage 2 model=========
-    model = LitModel(conf)
-    state = torch.load(args.stage2_checkpoint_path, map_location='cpu')
-    model.load_state_dict(state, strict=True)
-    model.ema_model.eval()
-    model.ema_model.to(args.device);
-    #=================================
-    
-    
-    #======Audio Input=========
-    if conf.infer_type.startswith('mfcc'):
-        # MFCC features
-        wav, sr = librosa.load(args.test_audio_path, sr=16000)
-        input_values = python_speech_features.mfcc(signal=wav, samplerate=sr, numcep=13, winlen=0.025, winstep=0.01)
-        d_mfcc_feat = python_speech_features.base.delta(input_values, 1)
-        d_mfcc_feat2 = python_speech_features.base.delta(input_values, 2)
-        audio_driven_obj = np.hstack((input_values, d_mfcc_feat, d_mfcc_feat2))
-        frame_start, frame_end = 0, int(audio_driven_obj.shape[0]/4)
-        audio_start, audio_end = int(frame_start * 4), int(frame_end * 4) # The video frame is fixed to 25 hz and the audio is fixed to 100 hz
-        
-        audio_driven = torch.Tensor(audio_driven_obj[audio_start:audio_end,:]).unsqueeze(0).float().to(args.device)
-        
-    elif conf.infer_type.startswith('hubert'):
-        # Hubert features
-        if not os.path.exists(args.test_hubert_path):
-            
-            if not check_package_installed('transformers'):
-                print('Please install transformers module first.')
-                exit(0)
-            hubert_model_path = 'ckpts/chinese-hubert-large'
-            if not os.path.exists(hubert_model_path):
-                print('Please download the hubert weight into the ckpts path first.')
-                exit(0)
-            print('You did not extract the audio features in advance, extracting online now, which will increase processing delay')
-
-            start_time = time.time()
-
-            # load hubert model
-            from transformers import Wav2Vec2FeatureExtractor, HubertModel
-            audio_model = HubertModel.from_pretrained(hubert_model_path).to(args.device)
-            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_path)
-            audio_model.feature_extractor._freeze_parameters()
-            audio_model.eval()
-
-            # hubert model forward pass
-            audio, sr = librosa.load(args.test_audio_path, sr=16000)
-            input_values = feature_extractor(audio, sampling_rate=16000, padding=True, do_normalize=True, return_tensors="pt").input_values
-            input_values = input_values.to(args.device)
-            ws_feats = []
-            with torch.no_grad():
-                outputs = audio_model(input_values, output_hidden_states=True)
-                for i in range(len(outputs.hidden_states)):
-                    ws_feats.append(outputs.hidden_states[i].detach().cpu().numpy())
-                ws_feat_obj = np.array(ws_feats)
-                ws_feat_obj = np.squeeze(ws_feat_obj, 1)
-                ws_feat_obj = np.pad(ws_feat_obj, ((0, 0), (0, 1), (0, 0)), 'edge') # align the audio length with video frame
-            
-            execution_time = time.time() - start_time
-            print(f"Extraction Audio Feature: {execution_time:.2f} Seconds")
-
-            audio_driven_obj = ws_feat_obj
-        else:
-            print(f'Using audio feature from path: {args.test_hubert_path}')
-            audio_driven_obj = np.load(args.test_hubert_path)
-
-        frame_start, frame_end = 0, int(audio_driven_obj.shape[1]/2)
-        audio_start, audio_end = int(frame_start * 2), int(frame_end * 2) # The video frame is fixed to 25 hz and the audio is fixed to 50 hz
-        
-        audio_driven = torch.Tensor(audio_driven_obj[:,audio_start:audio_end,:]).unsqueeze(0).float().to(args.device)
-    #============================
-    
-    # Diffusion Noise
-    noisyT = th.randn((1,frame_end, args.motion_dim)).to(args.device)
-    
-    #======Inputs for Attribute Control=========
-    if os.path.exists(args.pose_driven_path):
-        pose_obj = np.load(args.pose_driven_path)
-
-       
-        if len(pose_obj.shape) != 2:
-            print('please check your pose information. The shape must be like (T, 3).')
-            exit(0)
-        if pose_obj.shape[1] != 3:
-            print('please check your pose information. The shape must be like (T, 3).')
-            exit(0)
-
-        if pose_obj.shape[0] >= frame_end:
-            pose_obj = pose_obj[:frame_end,:]
-        else:
-            padding = np.tile(pose_obj[-1, :], (frame_end - pose_obj.shape[0], 1))
-            pose_obj = np.vstack((pose_obj, padding))
-            
-        pose_signal = torch.Tensor(pose_obj).unsqueeze(0).to(args.device) / 90 # 90 is for normalization here
-    else:
-        yaw_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_yaw
-        pitch_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_pitch
-        roll_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.pose_roll
-        pose_signal = torch.cat((yaw_signal, pitch_signal, roll_signal), dim=-1)
-    
-    pose_signal = torch.clamp(pose_signal, -1, 1)
-
-    face_location_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.face_location
-    face_scae_signal = torch.zeros(1, frame_end, 1).to(args.device) + args.face_scale
-    #===========================================
-
-    start_time = time.time()
-
-    #======Diffusion Denosing Process=========
-    generated_directions = model.render(one_shot_lia_start, one_shot_lia_direction, audio_driven, face_location_signal, face_scae_signal, pose_signal, noisyT, args.step_T, control_flag=args.control_flag)
-    #=========================================
-    
-    execution_time = time.time() - start_time
-    print(f"Motion Diffusion Model: {execution_time:.2f} Seconds")
-
-    generated_directions = generated_directions.detach().cpu().numpy()
-    
-    start_time = time.time()
-    #======Rendering images frame-by-frame=========
-    for pred_index in tqdm(range(generated_directions.shape[1])):
-        ori_img_recon = lia.render(one_shot_lia_start, torch.Tensor(generated_directions[:,pred_index,:]).to(args.device), feats)
-        ori_img_recon = ori_img_recon.clamp(-1, 1)
-        wav_pred = (ori_img_recon.detach() + 1) / 2
-        saved_image(wav_pred, os.path.join(frames_result_saved_path, "%06d.png"%(pred_index)))
-    #==============================================
-    
-    execution_time = time.time() - start_time
-    print(f"Renderer Model: {execution_time:.2f} Seconds")
-
-    frames_to_video(frames_result_saved_path, args.test_audio_path, predicted_video_256_path)
-    
-    shutil.rmtree(frames_result_saved_path)
-    
-    
-    # Enhancer
-    # Code is modified from https://github.com/OpenTalker/SadTalker/blob/cd4c0465ae0b54a6f85af57f5c65fec9fe23e7f8/src/utils/face_enhancer.py#L26
-
-    if args.face_sr and check_package_installed('gfpgan'):
-        from face_sr.face_enhancer import enhancer_list
-        import imageio
-
-        # Super-resolution
-        imageio.mimsave(predicted_video_512_path+'.tmp.mp4', enhancer_list(predicted_video_256_path, method='gfpgan', bg_upsampler=None), fps=float(25))
-        
-        # Merge audio and video
-        video_clip = VideoFileClip(predicted_video_512_path+'.tmp.mp4')
-        audio_clip = AudioFileClip(predicted_video_256_path)
-        final_clip = video_clip.set_audio(audio_clip)
-        final_clip.write_videofile(predicted_video_512_path, codec='libx264', audio_codec='aac')
-        
-        os.remove(predicted_video_512_path+'.tmp.mp4')
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--infer_type', type=str, default='mfcc_pose_only', help='mfcc_pose_only or mfcc_full_control')
-    parser.add_argument('--test_image_path', type=str, default='./test_demos/portraits/monalisa.jpg', help='Path to the portrait')
-    parser.add_argument('--test_audio_path', type=str, default='./test_demos/audios/english_female.wav', help='Path to the driven audio')
-    parser.add_argument('--test_hubert_path', type=str, default='./test_demos/audios_hubert/english_female.npy', help='Path to the driven audio(hubert type). Not needed for MFCC')
-    parser.add_argument('--result_path', type=str, default='./results/', help='Type of inference')
-    parser.add_argument('--stage1_checkpoint_path', type=str, default='./ckpts/stage1.ckpt', help='Path to the checkpoint of Stage1')
-    parser.add_argument('--stage2_checkpoint_path', type=str, default='./ckpts/pose_only.ckpt', help='Path to the checkpoint of Stage2')
-    parser.add_argument('--seed', type=int, default=0, help='seed for generations')
-    parser.add_argument('--control_flag', action='store_true', help='Whether to use control signal or not')
-    parser.add_argument('--pose_yaw', type=float, default=0.25, help='range from -1 to 1 (-90 ~ 90 angles)')
-    parser.add_argument('--pose_pitch', type=float, default=0, help='range from -1 to 1 (-90 ~ 90 angles)')
-    parser.add_argument('--pose_roll', type=float, default=0, help='range from -1 to 1 (-90 ~ 90 angles)')
-    parser.add_argument('--face_location', type=float, default=0.5, help='range from 0 to 1 (from left to right)')
-    parser.add_argument('--pose_driven_path', type=str, default='xxx', help='path to pose numpy, shape is (T, 3). You can check the following code https://github.com/liutaocode/talking_face_preprocessing to extract the yaw, pitch and roll.')
-    parser.add_argument('--face_scale', type=float, default=0.5, help='range from 0 to 1 (from small to large)')
-    parser.add_argument('--step_T', type=int, default=50, help='Step T for diffusion denoising process')
-    parser.add_argument('--image_size', type=int, default=256, help='Size of the image. Do not change.')
-    parser.add_argument('--device', type=str, default='cuda:0', help='Device for computation')
-    parser.add_argument('--motion_dim', type=int, default=20, help='Dimension of motion. Do not change.')
-    parser.add_argument('--decoder_layers', type=int, default=2, help='Layer number for the conformer.')
-    parser.add_argument('--face_sr', action='store_true', help='Face super-resolution (Optional). Please install GFPGAN first')
-
-
-
-    args = parser.parse_args()
-
-    main(args)

code/diffusion/__init__.py

@@ -1,6 +0,0 @@
-from typing import Union
-
-from .diffusion import SpacedDiffusionBeatGans, SpacedDiffusionBeatGansConfig
-
-Sampler = Union[SpacedDiffusionBeatGans]
-SamplerConfig = Union[SpacedDiffusionBeatGansConfig]

code/diffusion/base.py

@@ -1,1128 +0,0 @@
-"""
-This code started out as a PyTorch port of Ho et al's diffusion models:
-https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
-
-Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
-"""
-
-from model.unet_autoenc import AutoencReturn
-from config_base import BaseConfig
-import enum
-import math
-
-import numpy as np
-import torch as th
-from model import *
-from model.nn import mean_flat
-from typing import NamedTuple, Tuple
-from choices import *
-from torch.cuda.amp import autocast
-import torch.nn.functional as F
-
-from dataclasses import dataclass
-
-
-@dataclass
-class GaussianDiffusionBeatGansConfig(BaseConfig):
-    gen_type: GenerativeType
-    betas: Tuple[float]
-    model_type: ModelType
-    model_mean_type: ModelMeanType
-    model_var_type: ModelVarType
-    loss_type: LossType
-    rescale_timesteps: bool
-    fp16: bool
-    train_pred_xstart_detach: bool = True
-
-    def make_sampler(self):
-        return GaussianDiffusionBeatGans(self)
-
-
-class GaussianDiffusionBeatGans:
-    """
-    Utilities for training and sampling diffusion models.
-
-    Ported directly from here, and then adapted over time to further experimentation.
-    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
-
-    :param betas: a 1-D numpy array of betas for each diffusion timestep,
-                  starting at T and going to 1.
-    :param model_mean_type: a ModelMeanType determining what the model outputs.
-    :param model_var_type: a ModelVarType determining how variance is output.
-    :param loss_type: a LossType determining the loss function to use.
-    :param rescale_timesteps: if True, pass floating point timesteps into the
-                              model so that they are always scaled like in the
-                              original paper (0 to 1000).
-    """
-    def __init__(self, conf: GaussianDiffusionBeatGansConfig):
-        self.conf = conf
-        self.model_mean_type = conf.model_mean_type
-        self.model_var_type = conf.model_var_type
-        self.loss_type = conf.loss_type
-        self.rescale_timesteps = conf.rescale_timesteps
-
-        # Use float64 for accuracy.
-        betas = np.array(conf.betas, dtype=np.float64)
-        self.betas = betas
-        assert len(betas.shape) == 1, "betas must be 1-D"
-        assert (betas > 0).all() and (betas <= 1).all()
-
-        self.num_timesteps = int(betas.shape[0])
-
-        alphas = 1.0 - betas
-        self.alphas_cumprod = np.cumprod(alphas, axis=0)
-        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
-        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
-        assert self.alphas_cumprod_prev.shape == (self.num_timesteps, )
-
-        # calculations for diffusion q(x_t | x_{t-1}) and others
-        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
-        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
-        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
-        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
-        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod -
-                                                   1)
-
-        # calculations for posterior q(x_{t-1} | x_t, x_0)
-        self.posterior_variance = (betas * (1.0 - self.alphas_cumprod_prev) /
-                                   (1.0 - self.alphas_cumprod))
-        # log calculation clipped because the posterior variance is 0 at the
-        # beginning of the diffusion chain.
-        self.posterior_log_variance_clipped = np.log(
-            np.append(self.posterior_variance[1], self.posterior_variance[1:]))
-        self.posterior_mean_coef1 = (betas *
-                                     np.sqrt(self.alphas_cumprod_prev) /
-                                     (1.0 - self.alphas_cumprod))
-        self.posterior_mean_coef2 = ((1.0 - self.alphas_cumprod_prev) *
-                                     np.sqrt(alphas) /
-                                     (1.0 - self.alphas_cumprod))
-
-    def training_losses(self,
-                        model,
-                        motion_direction_start: th.Tensor,
-                        motion_target: th.Tensor,
-                        motion_start: th.Tensor,
-                        audio_feats: th.Tensor,
-                        face_location: th.Tensor,
-                        face_scale: th.Tensor,
-                        yaw_pitch_roll: th.Tensor,
-                        t: th.Tensor,
-                        model_kwargs=None,
-                        noise: th.Tensor = None):
-        """
-        Compute training losses for a single timestep.
-
-        :param model: the model to evaluate loss on.
-        :param x_start: the [N x C x ...] tensor of inputs.
-        :param t: a batch of timestep indices.
-        :param model_kwargs: if not None, a dict of extra keyword arguments to
-            pass to the model. This can be used for conditioning.
-        :param noise: if specified, the specific Gaussian noise to try to remove.
-        :return: a dict with the key "loss" containing a tensor of shape [N].
-                 Some mean or variance settings may also have other keys.
-        """
-        if model_kwargs is None:
-            model_kwargs = {}
-        if noise is None:
-            noise = th.randn_like(motion_target)
-
-        x_t = self.q_sample(motion_target, t, noise=noise)
-
-        terms = {'x_t': x_t}
-        
-        if self.loss_type in [
-                LossType.mse,
-                LossType.l1,
-        ]:
-            with autocast(self.conf.fp16):
-                # x_t is static wrt. to the diffusion process
-                predicted_direction, predicted_location, predicted_scale, predicted_pose = model.forward(motion_start, 
-                                              motion_direction_start,
-                                              audio_feats, 
-                                              face_location, 
-                                              face_scale,
-                                              yaw_pitch_roll,
-                                              x_t.detach(),
-                                              self._scale_timesteps(t),
-                                              control_flag=False)
-
-
-            target_types = {
-                ModelMeanType.eps: noise,
-            }
-            target = target_types[self.model_mean_type]
-            assert predicted_direction.shape == target.shape == motion_target.shape
-
-            if self.loss_type == LossType.mse:
-                if self.model_mean_type == ModelMeanType.eps:
-
-                    direction_loss = mean_flat((target - predicted_direction)**2)
-                    # import pdb;pdb.set_trace()
-                    location_loss = mean_flat((face_location.unsqueeze(-1) - predicted_location)**2)
-                    scale_loss = mean_flat((face_scale - predicted_scale)**2)
-                    pose_loss = mean_flat((yaw_pitch_roll - predicted_pose)**2)
-
-                    terms["mse"] = direction_loss + location_loss + scale_loss + pose_loss
-                   
-                else:
-                    raise NotImplementedError()
-            elif self.loss_type == LossType.l1:
-                # (n, c, h, w) => (n, )
-                terms["mse"] = mean_flat((target - predicted_direction).abs())
-            else:
-                raise NotImplementedError()
-
-            if "vb" in terms:
-                # if learning the variance also use the vlb loss
-                terms["loss"] = terms["mse"] + terms["vb"]
-            else:
-                terms["loss"] = terms["mse"]
-        else:
-            raise NotImplementedError(self.loss_type)
-        
-
-        return terms
-
-    def sample(self,
-               model: Model,
-               shape=None,
-               noise=None,
-               cond=None,
-               x_start=None,
-               clip_denoised=True,
-               model_kwargs=None,
-               progress=False):
-        """
-        Args:
-            x_start: given for the autoencoder
-        """
-        if model_kwargs is None:
-            model_kwargs = {}
-            if self.conf.model_type.has_autoenc():
-                model_kwargs['x_start'] = x_start
-                model_kwargs['cond'] = cond
-
-        if self.conf.gen_type == GenerativeType.ddpm:
-            return self.p_sample_loop(model,
-                                      shape=shape,
-                                      noise=noise,
-                                      clip_denoised=clip_denoised,
-                                      model_kwargs=model_kwargs,
-                                      progress=progress)
-        elif self.conf.gen_type == GenerativeType.ddim:
-            return self.ddim_sample_loop(model,
-                                         shape=shape,
-                                         noise=noise,
-                                         clip_denoised=clip_denoised,
-                                         model_kwargs=model_kwargs,
-                                         progress=progress)
-        else:
-            raise NotImplementedError()
-
-    def q_mean_variance(self, x_start, t):
-        """
-        Get the distribution q(x_t | x_0).
-
-        :param x_start: the [N x C x ...] tensor of noiseless inputs.
-        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
-        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
-        """
-        mean = (
-            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) *
-            x_start)
-        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t,
-                                        x_start.shape)
-        log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod,
-                                            t, x_start.shape)
-        return mean, variance, log_variance
-
-    def q_sample(self, x_start, t, noise=None):
-        """
-        Diffuse the data for a given number of diffusion steps.
-
-        In other words, sample from q(x_t | x_0).
-
-        :param x_start: the initial data batch.
-        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
-        :param noise: if specified, the split-out normal noise.
-        :return: A noisy version of x_start.
-        """
-        if noise is None:
-            noise = th.randn_like(x_start)
-        assert noise.shape == x_start.shape
-        return (
-            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) *
-            x_start + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod,
-                                           t, x_start.shape) * noise)
-
-    def q_posterior_mean_variance(self, x_start, x_t, t):
-        """
-        Compute the mean and variance of the diffusion posterior:
-
-            q(x_{t-1} | x_t, x_0)
-
-        """
-        assert x_start.shape == x_t.shape
-        posterior_mean = (
-            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) *
-            x_start +
-            _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) *
-            x_t)
-        posterior_variance = _extract_into_tensor(self.posterior_variance, t,
-                                                  x_t.shape)
-        posterior_log_variance_clipped = _extract_into_tensor(
-            self.posterior_log_variance_clipped, t, x_t.shape)
-        assert (posterior_mean.shape[0] == posterior_variance.shape[0] ==
-                posterior_log_variance_clipped.shape[0] == x_start.shape[0])
-        return posterior_mean, posterior_variance, posterior_log_variance_clipped
-
-    def p_mean_variance(self,
-                        model,
-                        x,
-                        t,
-                        clip_denoised=True,
-                        denoised_fn=None,
-                        model_kwargs=None):
-        """
-        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
-        the initial x, x_0.
-
-        :param model: the model, which takes a signal and a batch of timesteps
-                      as input.
-        :param x: the [N x C x ...] tensor at time t.
-        :param t: a 1-D Tensor of timesteps.
-        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
-        :param denoised_fn: if not None, a function which applies to the
-            x_start prediction before it is used to sample. Applies before
-            clip_denoised.
-        :param model_kwargs: if not None, a dict of extra keyword arguments to
-            pass to the model. This can be used for conditioning.
-        :return: a dict with the following keys:
-                 - 'mean': the model mean output.
-                 - 'variance': the model variance output.
-                 - 'log_variance': the log of 'variance'.
-                 - 'pred_xstart': the prediction for x_0.
-        """
-        if model_kwargs is None:
-            model_kwargs = {}
-
-        motion_start = model_kwargs['start']
-        audio_feats = model_kwargs['audio_driven']
-        face_location = model_kwargs['face_location']
-        face_scale = model_kwargs['face_scale']
-        yaw_pitch_roll = model_kwargs['yaw_pitch_roll']
-        motion_direction_start = model_kwargs['motion_direction_start']
-        control_flag = model_kwargs['control_flag']
-        
-        B, C = x.shape[:2]
-        assert t.shape == (B, )
-        with autocast(self.conf.fp16):
-            model_forward, _, _, _ = model.forward(motion_start, 
-                                                motion_direction_start,
-                                                audio_feats, 
-                                                face_location, 
-                                                face_scale,
-                                                yaw_pitch_roll,
-                                                x,
-                                                self._scale_timesteps(t), 
-                                                control_flag)
-        model_output = model_forward
-
-        if self.model_var_type in [
-                ModelVarType.fixed_large, ModelVarType.fixed_small
-        ]:
-            model_variance, model_log_variance = {
-                # for fixedlarge, we set the initial (log-)variance like so
-                # to get a better decoder log likelihood.
-                ModelVarType.fixed_large: (
-                    np.append(self.posterior_variance[1], self.betas[1:]),
-                    np.log(
-                        np.append(self.posterior_variance[1], self.betas[1:])),
-                ),
-                ModelVarType.fixed_small: (
-                    self.posterior_variance,
-                    self.posterior_log_variance_clipped,
-                ),
-            }[self.model_var_type]
-            model_variance = _extract_into_tensor(model_variance, t, x.shape)
-            model_log_variance = _extract_into_tensor(model_log_variance, t,
-                                                      x.shape)
-
-        def process_xstart(x):
-            if denoised_fn is not None:
-                x = denoised_fn(x)
-            if clip_denoised:
-                return x.clamp(-1, 1)
-            return x
-
-        if self.model_mean_type in [
-                ModelMeanType.eps,
-        ]:
-            if self.model_mean_type == ModelMeanType.eps:
-                pred_xstart = process_xstart(
-                    self._predict_xstart_from_eps(x_t=x, t=t,
-                                                  eps=model_output))
-            else:
-                raise NotImplementedError()
-            model_mean, _, _ = self.q_posterior_mean_variance(
-                x_start=pred_xstart, x_t=x, t=t)
-        else:
-            raise NotImplementedError(self.model_mean_type)
-
-        assert (model_mean.shape == model_log_variance.shape ==
-                pred_xstart.shape == x.shape)
-        return {
-            "mean": model_mean,
-            "variance": model_variance,
-            "log_variance": model_log_variance,
-            "pred_xstart": pred_xstart,
-            'model_forward': model_forward,
-        }
-
-    def _predict_xstart_from_eps(self, x_t, t, eps):
-        assert x_t.shape == eps.shape
-        return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t,
-                                     x_t.shape) * x_t -
-                _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,
-                                     x_t.shape) * eps)
-
-    def _predict_xstart_from_xprev(self, x_t, t, xprev):
-        assert x_t.shape == xprev.shape
-        return (  # (xprev - coef2*x_t) / coef1
-            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape)
-            * xprev - _extract_into_tensor(
-                self.posterior_mean_coef2 / self.posterior_mean_coef1, t,
-                x_t.shape) * x_t)
-
-    def _predict_xstart_from_scaled_xstart(self, t, scaled_xstart):
-        return scaled_xstart * _extract_into_tensor(
-            self.sqrt_recip_alphas_cumprod, t, scaled_xstart.shape)
-
-    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
-        return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t,
-                                     x_t.shape) * x_t -
-                pred_xstart) / _extract_into_tensor(
-                    self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
-
-    def _predict_eps_from_scaled_xstart(self, x_t, t, scaled_xstart):
-        """
-        Args:
-            scaled_xstart: is supposed to be sqrt(alphacum) * x_0
-        """
-        # 1 / sqrt(1-alphabar) * (x_t - scaled xstart)
-        return (x_t - scaled_xstart) / _extract_into_tensor(
-            self.sqrt_one_minus_alphas_cumprod, t, x_t.shape)
-
-    def _scale_timesteps(self, t):
-        if self.rescale_timesteps:
-            # scale t to be maxed out at 1000 steps
-            return t.float() * (1000.0 / self.num_timesteps)
-        return t
-
-    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
-        """
-        Compute the mean for the previous step, given a function cond_fn that
-        computes the gradient of a conditional log probability with respect to
-        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
-        condition on y.
-
-        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
-        """
-        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
-        new_mean = (p_mean_var["mean"].float() +
-                    p_mean_var["variance"] * gradient.float())
-        return new_mean
-
-    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
-        """
-        Compute what the p_mean_variance output would have been, should the
-        model's score function be conditioned by cond_fn.
-
-        See condition_mean() for details on cond_fn.
-
-        Unlike condition_mean(), this instead uses the conditioning strategy
-        from Song et al (2020).
-        """
-        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
-
-        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
-        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
-            x, self._scale_timesteps(t), **model_kwargs)
-
-        out = p_mean_var.copy()
-        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
-        out["mean"], _, _ = self.q_posterior_mean_variance(
-            x_start=out["pred_xstart"], x_t=x, t=t)
-        return out
-
-    def p_sample(
-        self,
-        model: Model,
-        x,
-        t,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-    ):
-        """
-        Sample x_{t-1} from the model at the given timestep.
-
-        :param model: the model to sample from.
-        :param x: the current tensor at x_{t-1}.
-        :param t: the value of t, starting at 0 for the first diffusion step.
-        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
-        :param denoised_fn: if not None, a function which applies to the
-            x_start prediction before it is used to sample.
-        :param cond_fn: if not None, this is a gradient function that acts
-                        similarly to the model.
-        :param model_kwargs: if not None, a dict of extra keyword arguments to
-            pass to the model. This can be used for conditioning.
-        :return: a dict containing the following keys:
-                 - 'sample': a random sample from the model.
-                 - 'pred_xstart': a prediction of x_0.
-        """
-        out = self.p_mean_variance(
-            model,
-            x,
-            t,
-            clip_denoised=clip_denoised,
-            denoised_fn=denoised_fn,
-            model_kwargs=model_kwargs,
-        )
-        noise = th.randn_like(x)
-        nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
-                        )  # no noise when t == 0
-        if cond_fn is not None:
-            out["mean"] = self.condition_mean(cond_fn,
-                                              out,
-                                              x,
-                                              t,
-                                              model_kwargs=model_kwargs)
-        sample = out["mean"] + nonzero_mask * th.exp(
-            0.5 * out["log_variance"]) * noise
-        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
-
-    def p_sample_loop(
-        self,
-        model: Model,
-        shape=None,
-        noise=None,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-        device=None,
-        progress=False,
-    ):
-        """
-        Generate samples from the model.
-
-        :param model: the model module.
-        :param shape: the shape of the samples, (N, C, H, W).
-        :param noise: if specified, the noise from the encoder to sample.
-                      Should be of the same shape as `shape`.
-        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
-        :param denoised_fn: if not None, a function which applies to the
-            x_start prediction before it is used to sample.
-        :param cond_fn: if not None, this is a gradient function that acts
-                        similarly to the model.
-        :param model_kwargs: if not None, a dict of extra keyword arguments to
-            pass to the model. This can be used for conditioning.
-        :param device: if specified, the device to create the samples on.
-                       If not specified, use a model parameter's device.
-        :param progress: if True, show a tqdm progress bar.
-        :return: a non-differentiable batch of samples.
-        """
-        final = None
-        for sample in self.p_sample_loop_progressive(
-                model,
-                shape,
-                noise=noise,
-                clip_denoised=clip_denoised,
-                denoised_fn=denoised_fn,
-                cond_fn=cond_fn,
-                model_kwargs=model_kwargs,
-                device=device,
-                progress=progress,
-        ):
-            final = sample
-        return final["sample"]
-
-    def p_sample_loop_progressive(
-        self,
-        model: Model,
-        shape=None,
-        noise=None,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-        device=None,
-        progress=False,
-    ):
-        """
-        Generate samples from the model and yield intermediate samples from
-        each timestep of diffusion.
-
-        Arguments are the same as p_sample_loop().
-        Returns a generator over dicts, where each dict is the return value of
-        p_sample().
-        """
-        if device is None:
-            device = next(model.parameters()).device
-        if noise is not None:
-            img = noise
-        else:
-            assert isinstance(shape, (tuple, list))
-            img = th.randn(*shape, device=device)
-        indices = list(range(self.num_timesteps))[::-1]
-
-        if progress:
-            # Lazy import so that we don't depend on tqdm.
-            from tqdm.auto import tqdm
-
-            indices = tqdm(indices)
-
-        for i in indices:
-            # t = th.tensor([i] * shape[0], device=device)
-            t = th.tensor([i] * len(img), device=device)
-            with th.no_grad():
-                out = self.p_sample(
-                    model,
-                    img,
-                    t,
-                    clip_denoised=clip_denoised,
-                    denoised_fn=denoised_fn,
-                    cond_fn=cond_fn,
-                    model_kwargs=model_kwargs,
-                )
-                yield out
-                img = out["sample"]
-
-    def ddim_sample(
-        self,
-        model: Model,
-        x,
-        t,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-        eta=0.0,
-    ):
-        """
-        Sample x_{t-1} from the model using DDIM.
-
-        Same usage as p_sample().
-        """
-        out = self.p_mean_variance(
-            model,
-            x,
-            t,
-            clip_denoised=clip_denoised,
-            denoised_fn=denoised_fn,
-            model_kwargs=model_kwargs,
-        )
-        if cond_fn is not None:
-            out = self.condition_score(cond_fn,
-                                       out,
-                                       x,
-                                       t,
-                                       model_kwargs=model_kwargs)
-
-        # Usually our model outputs epsilon, but we re-derive it
-        # in case we used x_start or x_prev prediction.
-        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
-
-        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
-        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t,
-                                              x.shape)
-        sigma = (eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) *
-                 th.sqrt(1 - alpha_bar / alpha_bar_prev))
-        # Equation 12.
-        noise = th.randn_like(x)
-        mean_pred = (out["pred_xstart"] * th.sqrt(alpha_bar_prev) +
-                     th.sqrt(1 - alpha_bar_prev - sigma**2) * eps)
-        nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
-                        )  # no noise when t == 0
-        sample = mean_pred + nonzero_mask * sigma * noise
-        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
-
-    def ddim_reverse_sample(
-        self,
-        model: Model,
-        x,
-        t,
-        clip_denoised=True,
-        denoised_fn=None,
-        model_kwargs=None,
-        eta=0.0,
-    ):
-        """
-        Sample x_{t+1} from the model using DDIM reverse ODE.
-        NOTE: never used ? 
-        """
-        assert eta == 0.0, "Reverse ODE only for deterministic path"
-        out = self.p_mean_variance(
-            model,
-            x,
-            t,
-            clip_denoised=clip_denoised,
-            denoised_fn=denoised_fn,
-            model_kwargs=model_kwargs,
-        )
-        # Usually our model outputs epsilon, but we re-derive it
-        # in case we used x_start or x_prev prediction.
-        eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape)
-               * x - out["pred_xstart"]) / _extract_into_tensor(
-                   self.sqrt_recipm1_alphas_cumprod, t, x.shape)
-        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t,
-                                              x.shape)
-
-        # Equation 12. reversed  (DDIM paper)  (th.sqrt == torch.sqrt)
-        mean_pred = (out["pred_xstart"] * th.sqrt(alpha_bar_next) +
-                     th.sqrt(1 - alpha_bar_next) * eps)
-
-        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
-
-    def ddim_reverse_sample_loop(
-        self,
-        model: Model,
-        x,
-        clip_denoised=True,
-        denoised_fn=None,
-        model_kwargs=None,
-        eta=0.0,
-        device=None,
-    ):
-        if device is None:
-            device = next(model.parameters()).device
-        sample_t = []
-        xstart_t = []
-        T = []
-        indices = list(range(self.num_timesteps))
-        sample = x
-        for i in indices:
-            t = th.tensor([i] * len(sample), device=device)
-            with th.no_grad():
-                out = self.ddim_reverse_sample(model,
-                                               sample,
-                                               t=t,
-                                               clip_denoised=clip_denoised,
-                                               denoised_fn=denoised_fn,
-                                               model_kwargs=model_kwargs,
-                                               eta=eta)
-                sample = out['sample']
-                # [1, ..., T]
-                sample_t.append(sample)
-                # [0, ...., T-1]
-                xstart_t.append(out['pred_xstart'])
-                # [0, ..., T-1] ready to use
-                T.append(t)
-
-        return {
-            #  xT "
-            'sample': sample,
-            # (1, ..., T)
-            'sample_t': sample_t,
-            # xstart here is a bit different from sampling from T = T-1 to T = 0
-            # may not be exact
-            'xstart_t': xstart_t,
-            'T': T,
-        }
-
-    def ddim_sample_loop(
-        self,
-        model: Model,
-        shape=None,
-        noise=None,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-        device=None,
-        progress=False,
-        eta=0.0,
-    ):
-        """
-        Generate samples from the model using DDIM.
-
-        Same usage as p_sample_loop().
-        """
-        final = None
-        for sample in self.ddim_sample_loop_progressive(
-                model,
-                shape,
-                noise=noise,
-                clip_denoised=clip_denoised,
-                denoised_fn=denoised_fn,
-                cond_fn=cond_fn,
-                model_kwargs=model_kwargs,
-                device=device,
-                progress=progress,
-                eta=eta,
-        ):
-            final = sample
-        return final["sample"]
-
-    def ddim_sample_loop_progressive(
-        self,
-        model: Model,
-        shape=None,
-        noise=None,
-        clip_denoised=True,
-        denoised_fn=None,
-        cond_fn=None,
-        model_kwargs=None,
-        device=None,
-        progress=False,
-        eta=0.0,
-    ):
-        """
-        Use DDIM to sample from the model and yield intermediate samples from
-        each timestep of DDIM.
-
-        Same usage as p_sample_loop_progressive().
-        """
-        if device is None:
-            device = next(model.parameters()).device
-        if noise is not None:
-            img = noise
-        else:
-            assert isinstance(shape, (tuple, list))
-            img = th.randn(*shape, device=device)
-        indices = list(range(self.num_timesteps))[::-1]
-
-        if progress:
-            # Lazy import so that we don't depend on tqdm.
-            from tqdm.auto import tqdm
-
-            indices = tqdm(indices)
-
-        for i in indices:
-
-            if isinstance(model_kwargs, list):
-                # index dependent model kwargs
-                # (T-1, ..., 0)
-                _kwargs = model_kwargs[i]
-            else:
-                _kwargs = model_kwargs
-
-            t = th.tensor([i] * len(img), device=device)
-            with th.no_grad():
-                out = self.ddim_sample(
-                    model,
-                    img,
-                    t,
-                    clip_denoised=clip_denoised,
-                    denoised_fn=denoised_fn,
-                    cond_fn=cond_fn,
-                    model_kwargs=_kwargs,
-                    eta=eta,
-                )
-                out['t'] = t
-                yield out
-                img = out["sample"]
-
-    def _vb_terms_bpd(self,
-                      model: Model,
-                      x_start,
-                      x_t,
-                      t,
-                      clip_denoised=True,
-                      model_kwargs=None):
-        """
-        Get a term for the variational lower-bound.
-
-        The resulting units are bits (rather than nats, as one might expect).
-        This allows for comparison to other papers.
-
-        :return: a dict with the following keys:
-                 - 'output': a shape [N] tensor of NLLs or KLs.
-                 - 'pred_xstart': the x_0 predictions.
-        """
-        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
-            x_start=x_start, x_t=x_t, t=t)
-        out = self.p_mean_variance(model,
-                                   x_t,
-                                   t,
-                                   clip_denoised=clip_denoised,
-                                   model_kwargs=model_kwargs)
-        kl = normal_kl(true_mean, true_log_variance_clipped, out["mean"],
-                       out["log_variance"])
-        kl = mean_flat(kl) / np.log(2.0)
-
-        decoder_nll = -discretized_gaussian_log_likelihood(
-            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"])
-        assert decoder_nll.shape == x_start.shape
-        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
-
-        # At the first timestep return the decoder NLL,
-        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
-        output = th.where((t == 0), decoder_nll, kl)
-        return {
-            "output": output,
-            "pred_xstart": out["pred_xstart"],
-            'model_forward': out['model_forward'],
-        }
-
-    def _prior_bpd(self, x_start):
-        """
-        Get the prior KL term for the variational lower-bound, measured in
-        bits-per-dim.
-
-        This term can't be optimized, as it only depends on the encoder.
-
-        :param x_start: the [N x C x ...] tensor of inputs.
-        :return: a batch of [N] KL values (in bits), one per batch element.
-        """
-        batch_size = x_start.shape[0]
-        t = th.tensor([self.num_timesteps - 1] * batch_size,
-                      device=x_start.device)
-        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
-        kl_prior = normal_kl(mean1=qt_mean,
-                             logvar1=qt_log_variance,
-                             mean2=0.0,
-                             logvar2=0.0)
-        return mean_flat(kl_prior) / np.log(2.0)
-
-    def calc_bpd_loop(self,
-                      model: Model,
-                      x_start,
-                      clip_denoised=True,
-                      model_kwargs=None):
-        """
-        Compute the entire variational lower-bound, measured in bits-per-dim,
-        as well as other related quantities.
-
-        :param model: the model to evaluate loss on.
-        :param x_start: the [N x C x ...] tensor of inputs.
-        :param clip_denoised: if True, clip denoised samples.
-        :param model_kwargs: if not None, a dict of extra keyword arguments to
-            pass to the model. This can be used for conditioning.
-
-        :return: a dict containing the following keys:
-                 - total_bpd: the total variational lower-bound, per batch element.
-                 - prior_bpd: the prior term in the lower-bound.
-                 - vb: an [N x T] tensor of terms in the lower-bound.
-                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
-                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
-        """
-        device = x_start.device
-        batch_size = x_start.shape[0]
-
-        vb = []
-        xstart_mse = []
-        mse = []
-        for t in list(range(self.num_timesteps))[::-1]:
-            t_batch = th.tensor([t] * batch_size, device=device)
-            noise = th.randn_like(x_start)
-            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
-            # Calculate VLB term at the current timestep
-            with th.no_grad():
-                out = self._vb_terms_bpd(
-                    model,
-                    x_start=x_start,
-                    x_t=x_t,
-                    t=t_batch,
-                    clip_denoised=clip_denoised,
-                    model_kwargs=model_kwargs,
-                )
-            vb.append(out["output"])
-            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start)**2))
-            eps = self._predict_eps_from_xstart(x_t, t_batch,
-                                                out["pred_xstart"])
-            mse.append(mean_flat((eps - noise)**2))
-
-        vb = th.stack(vb, dim=1)
-        xstart_mse = th.stack(xstart_mse, dim=1)
-        mse = th.stack(mse, dim=1)
-
-        prior_bpd = self._prior_bpd(x_start)
-        total_bpd = vb.sum(dim=1) + prior_bpd
-        return {
-            "total_bpd": total_bpd,
-            "prior_bpd": prior_bpd,
-            "vb": vb,
-            "xstart_mse": xstart_mse,
-            "mse": mse,
-        }
-
-
-def _extract_into_tensor(arr, timesteps, broadcast_shape):
-    """
-    Extract values from a 1-D numpy array for a batch of indices.
-
-    :param arr: the 1-D numpy array.
-    :param timesteps: a tensor of indices into the array to extract.
-    :param broadcast_shape: a larger shape of K dimensions with the batch
-                            dimension equal to the length of timesteps.
-    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
-    """
-    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
-    while len(res.shape) < len(broadcast_shape):
-        res = res[..., None]
-    return res.expand(broadcast_shape)
-
-
-def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
-    """
-    Get a pre-defined beta schedule for the given name.
-
-    The beta schedule library consists of beta schedules which remain similar
-    in the limit of num_diffusion_timesteps.
-    Beta schedules may be added, but should not be removed or changed once
-    they are committed to maintain backwards compatibility.
-    """
-    if schedule_name == "linear":
-        # Linear schedule from Ho et al, extended to work for any number of
-        # diffusion steps.
-        scale = 1000 / num_diffusion_timesteps
-        beta_start = scale * 0.0001
-        beta_end = scale * 0.02
-        return np.linspace(beta_start,
-                           beta_end,
-                           num_diffusion_timesteps,
-                           dtype=np.float64)
-    elif schedule_name == "cosine":
-        return betas_for_alpha_bar(
-            num_diffusion_timesteps,
-            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2)**2,
-        )
-    elif schedule_name == "const0.01":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.01] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.015":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.015] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.008":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.008] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0065":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0065] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0055":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0055] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0045":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0045] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0035":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0035] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0025":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0025] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    elif schedule_name == "const0.0015":
-        scale = 1000 / num_diffusion_timesteps
-        return np.array([scale * 0.0015] * num_diffusion_timesteps,
-                        dtype=np.float64)
-    else:
-        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
-
-
-def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
-    """
-    Create a beta schedule that discretizes the given alpha_t_bar function,
-    which defines the cumulative product of (1-beta) over time from t = [0,1].
-
-    :param num_diffusion_timesteps: the number of betas to produce.
-    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
-                      produces the cumulative product of (1-beta) up to that
-                      part of the diffusion process.
-    :param max_beta: the maximum beta to use; use values lower than 1 to
-                     prevent singularities.
-    """
-    betas = []
-    for i in range(num_diffusion_timesteps):
-        t1 = i / num_diffusion_timesteps
-        t2 = (i + 1) / num_diffusion_timesteps
-        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
-    return np.array(betas)
-
-
-def normal_kl(mean1, logvar1, mean2, logvar2):
-    """
-    Compute the KL divergence between two gaussians.
-
-    Shapes are automatically broadcasted, so batches can be compared to
-    scalars, among other use cases.
-    """
-    tensor = None
-    for obj in (mean1, logvar1, mean2, logvar2):
-        if isinstance(obj, th.Tensor):
-            tensor = obj
-            break
-    assert tensor is not None, "at least one argument must be a Tensor"
-
-    # Force variances to be Tensors. Broadcasting helps convert scalars to
-    # Tensors, but it does not work for th.exp().
-    logvar1, logvar2 = [
-        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
-        for x in (logvar1, logvar2)
-    ]
-
-    return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) +
-                  ((mean1 - mean2)**2) * th.exp(-logvar2))
-
-
-def approx_standard_normal_cdf(x):
-    """
-    A fast approximation of the cumulative distribution function of the
-    standard normal.
-    """
-    return 0.5 * (
-        1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
-
-
-def discretized_gaussian_log_likelihood(x, *, means, log_scales):
-    """
-    Compute the log-likelihood of a Gaussian distribution discretizing to a
-    given image.
-
-    :param x: the target images. It is assumed that this was uint8 values,
-              rescaled to the range [-1, 1].
-    :param means: the Gaussian mean Tensor.
-    :param log_scales: the Gaussian log stddev Tensor.
-    :return: a tensor like x of log probabilities (in nats).
-    """
-    assert x.shape == means.shape == log_scales.shape
-    centered_x = x - means
-    inv_stdv = th.exp(-log_scales)
-    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
-    cdf_plus = approx_standard_normal_cdf(plus_in)
-    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
-    cdf_min = approx_standard_normal_cdf(min_in)
-    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
-    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
-    cdf_delta = cdf_plus - cdf_min
-    log_probs = th.where(
-        x < -0.999,
-        log_cdf_plus,
-        th.where(x > 0.999, log_one_minus_cdf_min,
-                 th.log(cdf_delta.clamp(min=1e-12))),
-    )
-    assert log_probs.shape == x.shape
-    return log_probs
-
-
-class DummyModel(th.nn.Module):
-    def __init__(self, pred):
-        super().__init__()
-        self.pred = pred
-
-    def forward(self, *args, **kwargs):
-        return DummyReturn(pred=self.pred)
-
-
-class DummyReturn(NamedTuple):
-    pred: th.Tensor

code/diffusion/diffusion.py

@@ -1,156 +0,0 @@
-from .base import *
-from dataclasses import dataclass
-
-
-def space_timesteps(num_timesteps, section_counts):
-    """
-    Create a list of timesteps to use from an original diffusion process,
-    given the number of timesteps we want to take from equally-sized portions
-    of the original process.
-
-    For example, if there's 300 timesteps and the section counts are [10,15,20]
-    then the first 100 timesteps are strided to be 10 timesteps, the second 100
-    are strided to be 15 timesteps, and the final 100 are strided to be 20.
-
-    If the stride is a string starting with "ddim", then the fixed striding
-    from the DDIM paper is used, and only one section is allowed.
-
-    :param num_timesteps: the number of diffusion steps in the original
-                          process to divide up.
-    :param section_counts: either a list of numbers, or a string containing
-                           comma-separated numbers, indicating the step count
-                           per section. As a special case, use "ddimN" where N
-                           is a number of steps to use the striding from the
-                           DDIM paper.
-    :return: a set of diffusion steps from the original process to use.
-    """
-    if isinstance(section_counts, str):
-        if section_counts.startswith("ddim"):
-            desired_count = int(section_counts[len("ddim"):])
-            for i in range(1, num_timesteps):
-                if len(range(0, num_timesteps, i)) == desired_count:
-                    return set(range(0, num_timesteps, i))
-            raise ValueError(
-                f"cannot create exactly {num_timesteps} steps with an integer stride"
-            )
-        section_counts = [int(x) for x in section_counts.split(",")]
-    size_per = num_timesteps // len(section_counts)
-    extra = num_timesteps % len(section_counts)
-    start_idx = 0
-    all_steps = []
-    for i, section_count in enumerate(section_counts):
-        size = size_per + (1 if i < extra else 0)
-        if size < section_count:
-            raise ValueError(
-                f"cannot divide section of {size} steps into {section_count}")
-        if section_count <= 1:
-            frac_stride = 1
-        else:
-            frac_stride = (size - 1) / (section_count - 1)
-        cur_idx = 0.0
-        taken_steps = []
-        for _ in range(section_count):
-            taken_steps.append(start_idx + round(cur_idx))
-            cur_idx += frac_stride
-        all_steps += taken_steps
-        start_idx += size
-    return set(all_steps)
-
-
-@dataclass
-class SpacedDiffusionBeatGansConfig(GaussianDiffusionBeatGansConfig):
-    use_timesteps: Tuple[int] = None
-
-    def make_sampler(self):
-        return SpacedDiffusionBeatGans(self)
-
-
-class SpacedDiffusionBeatGans(GaussianDiffusionBeatGans):
-    """
-    A diffusion process which can skip steps in a base diffusion process.
-
-    :param use_timesteps: a collection (sequence or set) of timesteps from the
-                          original diffusion process to retain.
-    :param kwargs: the kwargs to create the base diffusion process.
-    """
-    def __init__(self, conf: SpacedDiffusionBeatGansConfig):
-        self.conf = conf
-        self.use_timesteps = set(conf.use_timesteps)
-        # how the new t's mapped to the old t's
-        self.timestep_map = []
-        self.original_num_steps = len(conf.betas)
-
-        base_diffusion = GaussianDiffusionBeatGans(conf)  # pylint: disable=missing-kwoa
-        last_alpha_cumprod = 1.0
-        new_betas = []
-        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
-            if i in self.use_timesteps:
-                # getting the new betas of the new timesteps
-                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
-                last_alpha_cumprod = alpha_cumprod
-                self.timestep_map.append(i)
-        conf.betas = np.array(new_betas)
-        super().__init__(conf)
-
-    def p_mean_variance(self, model: Model, *args, **kwargs):  # pylint: disable=signature-differs
-        return super().p_mean_variance(self._wrap_model(model), *args,
-                                       **kwargs)
-
-    def training_losses(self, model: Model, *args, **kwargs):  # pylint: disable=signature-differs
-        return super().training_losses(self._wrap_model(model), *args,
-                                       **kwargs)
-
-    def condition_mean(self, cond_fn, *args, **kwargs):
-        return super().condition_mean(self._wrap_model(cond_fn), *args,
-                                      **kwargs)
-
-    def condition_score(self, cond_fn, *args, **kwargs):
-        return super().condition_score(self._wrap_model(cond_fn), *args,
-                                       **kwargs)
-
-    def _wrap_model(self, model: Model):
-        if isinstance(model, _WrappedModel):
-            return model
-        return _WrappedModel(model, self.timestep_map, self.rescale_timesteps,
-                             self.original_num_steps)
-
-    def _scale_timesteps(self, t):
-        # Scaling is done by the wrapped model.
-        return t
-
-
-class _WrappedModel:
-    """
-    converting the supplied t's to the old t's scales.
-    """
-    def __init__(self, model, timestep_map, rescale_timesteps,
-                 original_num_steps):
-        self.model = model
-        self.timestep_map = timestep_map
-        self.rescale_timesteps = rescale_timesteps
-        self.original_num_steps = original_num_steps
-
-    def forward(self,motion_start, motion_direction_start, audio_feats,face_location, face_scale,yaw_pitch_roll, x_t, t, control_flag=False):
-        """
-        Args:
-            t: t's with differrent ranges (can be << T due to smaller eval T) need to be converted to the original t's
-            t_cond: the same as t but can be of different values
-        """
-        map_tensor = th.tensor(self.timestep_map,
-                               device=t.device,
-                               dtype=t.dtype)
-
-        def do(t):
-            new_ts = map_tensor[t]
-            if self.rescale_timesteps:
-                new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
-            return new_ts
-
-        return self.model(motion_start, motion_direction_start, audio_feats,face_location, face_scale,yaw_pitch_roll, x_t,do(t), control_flag=control_flag)
-
-    def __getattr__(self, name):
-        # allow for calling the model's methods
-        if hasattr(self.model, name):
-            func = getattr(self.model, name)
-            return func
-        raise AttributeError(name)

code/diffusion/resample.py

@@ -1,63 +0,0 @@
-from abc import ABC, abstractmethod
-
-import numpy as np
-import torch as th
-import torch.distributed as dist
-
-
-def create_named_schedule_sampler(name, diffusion):
-    """
-    Create a ScheduleSampler from a library of pre-defined samplers.
-
-    :param name: the name of the sampler.
-    :param diffusion: the diffusion object to sample for.
-    """
-    if name == "uniform":
-        return UniformSampler(diffusion)
-    else:
-        raise NotImplementedError(f"unknown schedule sampler: {name}")
-
-
-class ScheduleSampler(ABC):
-    """
-    A distribution over timesteps in the diffusion process, intended to reduce
-    variance of the objective.
-
-    By default, samplers perform unbiased importance sampling, in which the
-    objective's mean is unchanged.
-    However, subclasses may override sample() to change how the resampled
-    terms are reweighted, allowing for actual changes in the objective.
-    """
-    @abstractmethod
-    def weights(self):
-        """
-        Get a numpy array of weights, one per diffusion step.
-
-        The weights needn't be normalized, but must be positive.
-        """
-
-    def sample(self, batch_size, device):
-        """
-        Importance-sample timesteps for a batch.
-
-        :param batch_size: the number of timesteps.
-        :param device: the torch device to save to.
-        :return: a tuple (timesteps, weights):
-                 - timesteps: a tensor of timestep indices.
-                 - weights: a tensor of weights to scale the resulting losses.
-        """
-        w = self.weights()
-        p = w / np.sum(w)
-        indices_np = np.random.choice(len(p), size=(batch_size, ), p=p)
-        indices = th.from_numpy(indices_np).long().to(device)
-        weights_np = 1 / (len(p) * p[indices_np])
-        weights = th.from_numpy(weights_np).float().to(device)
-        return indices, weights
-
-
-class UniformSampler(ScheduleSampler):
-    def __init__(self, num_timesteps):
-        self._weights = np.ones([num_timesteps])
-
-    def weights(self):
-        return self._weights

code/dist_utils.py

@@ -1,42 +0,0 @@
-from typing import List
-from torch import distributed
-
-
-def barrier():
-    if distributed.is_initialized():
-        distributed.barrier()
-    else:
-        pass
-
-
-def broadcast(data, src):
-    if distributed.is_initialized():
-        distributed.broadcast(data, src)
-    else:
-        pass
-
-
-def all_gather(data: List, src):
-    if distributed.is_initialized():
-        distributed.all_gather(data, src)
-    else:
-        data[0] = src
-
-
-def get_rank():
-    if distributed.is_initialized():
-        return distributed.get_rank()
-    else:
-        return 0
-
-
-def get_world_size():
-    if distributed.is_initialized():
-        return distributed.get_world_size()
-    else:
-        return 1
-
-
-def chunk_size(size, rank, world_size):
-    extra = rank < size % world_size
-    return size // world_size + extra

code/experiment.py

@@ -1,356 +0,0 @@
-import copy
-import os
-
-import numpy as np
-import pytorch_lightning as pl
-import torch
-from pytorch_lightning import loggers as pl_loggers
-from pytorch_lightning.callbacks import *
-from torch.cuda import amp
-from torch.optim.optimizer import Optimizer
-from torch.utils.data.dataset import TensorDataset
-from model.seq2seq import DiffusionPredictor
-
-from config import *
-from dist_utils import *
-from renderer import *
-
-# This part is modified from: https://github.com/phizaz/diffae/blob/master/experiment.py
-class LitModel(pl.LightningModule):
-    def __init__(self, conf: TrainConfig):
-        super().__init__()
-        assert conf.train_mode != TrainMode.manipulate
-        if conf.seed is not None:
-            pl.seed_everything(conf.seed)
-
-        self.save_hyperparameters(conf.as_dict_jsonable())
-
-        self.conf = conf
-
-        self.model = DiffusionPredictor(conf)
-        
-        self.ema_model = copy.deepcopy(self.model)
-        self.ema_model.requires_grad_(False)
-        self.ema_model.eval()
-
-        self.sampler = conf.make_diffusion_conf().make_sampler()
-        self.eval_sampler = conf.make_eval_diffusion_conf().make_sampler()
-
-        # this is shared for both model and latent
-        self.T_sampler = conf.make_T_sampler()
-
-        if conf.train_mode.use_latent_net():
-            self.latent_sampler = conf.make_latent_diffusion_conf(
-            ).make_sampler()
-            self.eval_latent_sampler = conf.make_latent_eval_diffusion_conf(
-            ).make_sampler()
-        else:
-            self.latent_sampler = None
-            self.eval_latent_sampler = None
-
-        # initial variables for consistent sampling
-        self.register_buffer(
-            'x_T',
-            torch.randn(conf.sample_size, 3, conf.img_size, conf.img_size))
-
-
-    def render(self, start, motion_direction_start, audio_driven, face_location, face_scale, ypr_info, noisyT, step_T, control_flag):
-        if step_T is None:
-            sampler = self.eval_sampler
-        else:
-            sampler = self.conf._make_diffusion_conf(step_T).make_sampler()
-
-        pred_img = render_condition(self.conf,
-                                        self.ema_model,
-                                        sampler, start, motion_direction_start, audio_driven, face_location, face_scale, ypr_info, noisyT, control_flag)
-        return pred_img
-
-    def forward(self, noise=None, x_start=None, ema_model: bool = False):
-        with amp.autocast(False):
-            if not self.disable_ema:
-                model = self.ema_model
-            else:
-                model = self.model
-            gen = self.eval_sampler.sample(model=model,
-                                           noise=noise,
-                                           x_start=x_start)
-            return gen
-
-    def setup(self, stage=None) -> None:
-        """
-        make datasets & seeding each worker separately
-        """
-        ##############################################
-        # NEED TO SET THE SEED SEPARATELY HERE
-        if self.conf.seed is not None:
-            seed = self.conf.seed * get_world_size() + self.global_rank
-            np.random.seed(seed)
-            torch.manual_seed(seed)
-            torch.cuda.manual_seed(seed)
-            print('local seed:', seed)
-        ##############################################
-
-        self.train_data = self.conf.make_dataset()
-        print('train data:', len(self.train_data))
-        self.val_data = self.train_data
-        print('val data:', len(self.val_data))
-
-    def _train_dataloader(self, drop_last=True):
-        """
-        really make the dataloader
-        """
-        # make sure to use the fraction of batch size
-        # the batch size is global!
-        conf = self.conf.clone()
-        conf.batch_size = self.batch_size
-
-        dataloader = conf.make_loader(self.train_data,
-                                      shuffle=True,
-                                      drop_last=drop_last)
-        return dataloader
-
-    def train_dataloader(self):
-        """
-        return the dataloader, if diffusion mode => return image dataset
-        if latent mode => return the inferred latent dataset
-        """
-        print('on train dataloader start ...')
-        if self.conf.train_mode.require_dataset_infer():
-            if self.conds is None:
-                # usually we load self.conds from a file
-                # so we do not need to do this again!
-                self.conds = self.infer_whole_dataset()
-                # need to use float32! unless the mean & std will be off!
-                # (1, c)
-                self.conds_mean.data = self.conds.float().mean(dim=0,
-                                                               keepdim=True)
-                self.conds_std.data = self.conds.float().std(dim=0,
-                                                             keepdim=True)
-            print('mean:', self.conds_mean.mean(), 'std:',
-                  self.conds_std.mean())
-
-            # return the dataset with pre-calculated conds
-            conf = self.conf.clone()
-            conf.batch_size = self.batch_size
-            data = TensorDataset(self.conds)
-            return conf.make_loader(data, shuffle=True)
-        else:
-            return self._train_dataloader()
-
-    @property
-    def batch_size(self):
-        """
-        local batch size for each worker
-        """
-        ws = get_world_size()
-        assert self.conf.batch_size % ws == 0
-        return self.conf.batch_size // ws
-
-    @property
-    def num_samples(self):
-        """
-        (global) batch size * iterations
-        """
-        # batch size here is global!
-        # global_step already takes into account the accum batches
-        return self.global_step * self.conf.batch_size_effective
-
-    def is_last_accum(self, batch_idx):
-        """
-        is it the last gradient accumulation loop? 
-        used with gradient_accum > 1 and to see if the optimizer will perform "step" in this iteration or not
-        """
-        return (batch_idx + 1) % self.conf.accum_batches == 0
-
-    def training_step(self, batch, batch_idx):
-        """
-        given an input, calculate the loss function
-        no optimization at this stage.
-        """
-        with amp.autocast(False):
-            motion_start = batch['motion_start'] # torch.Size([B, 512])
-            motion_direction = batch['motion_direction'] # torch.Size([B, 125, 20])
-            audio_feats = batch['audio_feats'].float() # torch.Size([B, 25, 250, 1024])
-            face_location = batch['face_location'].float() # torch.Size([B, 125])
-            face_scale = batch['face_scale'].float() # torch.Size([B, 125, 1])
-            yaw_pitch_roll = batch['yaw_pitch_roll'].float() # torch.Size([B, 125, 3])
-            motion_direction_start = batch['motion_direction_start'].float() # torch.Size([B, 20])
-
-            # import pdb; pdb.set_trace()
-            if self.conf.train_mode == TrainMode.diffusion:
-                """
-                main training mode!!!
-                """
-                # with numpy seed we have the problem that the sample t's are related!
-                t, weight = self.T_sampler.sample(len(motion_start), motion_start.device)
-                losses = self.sampler.training_losses(model=self.model,
-                                                      motion_direction_start=motion_direction_start,
-                                                      motion_target=motion_direction,
-                                                      motion_start=motion_start, 
-                                                      audio_feats=audio_feats, 
-                                                      face_location=face_location,
-                                                      face_scale=face_scale,
-                                                      yaw_pitch_roll=yaw_pitch_roll,
-                                                      t=t)
-            else:
-                raise NotImplementedError()
-
-            loss = losses['loss'].mean()
-            # divide by accum batches to make the accumulated gradient exact!
-            for key in losses.keys():
-                losses[key] = self.all_gather(losses[key]).mean()
-
-            if self.global_rank == 0:
-                self.logger.experiment.add_scalar('loss', losses['loss'],
-                                                  self.num_samples)
-                for key in losses:
-                    self.logger.experiment.add_scalar(
-                        f'loss/{key}', losses[key], self.num_samples)
-
-        return {'loss': loss}
-
-    def on_train_batch_end(self, outputs, batch, batch_idx: int,
-                           dataloader_idx: int) -> None:
-        """
-        after each training step ...
-        """
-        if self.is_last_accum(batch_idx):
-
-            if self.conf.train_mode == TrainMode.latent_diffusion:
-                # it trains only the latent hence change only the latent
-                ema(self.model.latent_net, self.ema_model.latent_net,
-                    self.conf.ema_decay)
-            else:
-                ema(self.model, self.ema_model, self.conf.ema_decay)
-
-    def on_before_optimizer_step(self, optimizer: Optimizer,
-                                 optimizer_idx: int) -> None:
-        # fix the fp16 + clip grad norm problem with pytorch lightinng
-        # this is the currently correct way to do it
-        if self.conf.grad_clip > 0:
-            # from trainer.params_grads import grads_norm, iter_opt_params
-            params = [
-                p for group in optimizer.param_groups for p in group['params']
-            ]
-            torch.nn.utils.clip_grad_norm_(params,
-                                           max_norm=self.conf.grad_clip)
-    def configure_optimizers(self):
-        out = {}
-        if self.conf.optimizer == OptimizerType.adam:
-            optim = torch.optim.Adam(self.model.parameters(),
-                                     lr=self.conf.lr,
-                                     weight_decay=self.conf.weight_decay)
-        elif self.conf.optimizer == OptimizerType.adamw:
-            optim = torch.optim.AdamW(self.model.parameters(),
-                                      lr=self.conf.lr,
-                                      weight_decay=self.conf.weight_decay)
-        else:
-            raise NotImplementedError()
-        out['optimizer'] = optim
-        if self.conf.warmup > 0:
-            sched = torch.optim.lr_scheduler.LambdaLR(optim,
-                                                      lr_lambda=WarmupLR(
-                                                          self.conf.warmup))
-            out['lr_scheduler'] = {
-                'scheduler': sched,
-                'interval': 'step',
-            }
-        return out
-
-    def split_tensor(self, x):
-        """
-        extract the tensor for a corresponding "worker" in the batch dimension
-
-        Args:
-            x: (n, c)
-
-        Returns: x: (n_local, c)
-        """
-        n = len(x)
-        rank = self.global_rank
-        world_size = get_world_size()
-        # print(f'rank: {rank}/{world_size}')
-        per_rank = n // world_size
-        return x[rank * per_rank:(rank + 1) * per_rank]
-
-def ema(source, target, decay):
-    source_dict = source.state_dict()
-    target_dict = target.state_dict()
-    for key in source_dict.keys():
-        target_dict[key].data.copy_(target_dict[key].data * decay +
-                                    source_dict[key].data * (1 - decay))
-
-
-class WarmupLR:
-    def __init__(self, warmup) -> None:
-        self.warmup = warmup
-
-    def __call__(self, step):
-        return min(step, self.warmup) / self.warmup
-
-
-def is_time(num_samples, every, step_size):
-    closest = (num_samples // every) * every
-    return num_samples - closest < step_size
-
-
-def train(conf: TrainConfig, gpus, nodes=1, mode: str = 'train'):
-    print('conf:', conf.name)
-    # assert not (conf.fp16 and conf.grad_clip > 0
-    #             ), 'pytorch lightning has bug with amp + gradient clipping'
-    model = LitModel(conf)
-
-    if not os.path.exists(conf.logdir):
-        os.makedirs(conf.logdir)
-    checkpoint = ModelCheckpoint(dirpath=f'{conf.logdir}',
-                                 save_last=True,
-                                 save_top_k=-1,
-                                 every_n_epochs=10)
-    checkpoint_path = f'{conf.logdir}/last.ckpt'
-    print('ckpt path:', checkpoint_path)
-    if os.path.exists(checkpoint_path):
-        resume = checkpoint_path
-        print('resume!')
-    else:
-        if conf.continue_from is not None:
-            # continue from a checkpoint
-            resume = conf.continue_from.pathcd 
-        else:
-            resume = None
-
-    tb_logger = pl_loggers.TensorBoardLogger(save_dir=conf.logdir,
-                                             name=None,
-                                             version='')
-
-    # from pytorch_lightning.
-
-    plugins = []
-    if len(gpus) == 1 and nodes == 1:
-        accelerator = None
-    else:
-        accelerator = 'ddp'
-        from pytorch_lightning.plugins import DDPPlugin
-
-        # important for working with gradient checkpoint
-        plugins.append(DDPPlugin(find_unused_parameters=True))
-
-    trainer = pl.Trainer(
-        max_steps=conf.total_samples // conf.batch_size_effective,
-        resume_from_checkpoint=resume,
-        gpus=gpus,
-        num_nodes=nodes,
-        accelerator=accelerator,
-        precision=16 if conf.fp16 else 32,
-        callbacks=[
-            checkpoint,
-            LearningRateMonitor(),
-        ],
-        # clip in the model instead
-        # gradient_clip_val=conf.grad_clip,
-        replace_sampler_ddp=True,
-        logger=tb_logger,
-        accumulate_grad_batches=conf.accum_batches,
-        plugins=plugins,
-    )
-
-    trainer.fit(model)

code/face_sr/face_enhancer.py

@@ -1,123 +0,0 @@
-import os
-import torch 
-
-from gfpgan import GFPGANer
-
-from tqdm import tqdm
-
-from .videoio import load_video_to_cv2
-
-import cv2
-
-
-class GeneratorWithLen(object):
-    """ From https://stackoverflow.com/a/7460929 """
-
-    def __init__(self, gen, length):
-        self.gen = gen
-        self.length = length
-
-    def __len__(self):
-        return self.length
-
-    def __iter__(self):
-        return self.gen
-
-def enhancer_list(images, method='gfpgan', bg_upsampler='realesrgan'):
-    gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler)
-    return list(gen)
-
-def enhancer_generator_with_len(images, method='gfpgan', bg_upsampler='realesrgan'):
-    """ Provide a generator with a __len__ method so that it can passed to functions that
-    call len()"""
-
-    if os.path.isfile(images): # handle video to images
-        # TODO: Create a generator version of load_video_to_cv2
-        images = load_video_to_cv2(images)
-
-    gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler)
-    gen_with_len = GeneratorWithLen(gen, len(images))
-    return gen_with_len
-
-def enhancer_generator_no_len(images, method='gfpgan', bg_upsampler='realesrgan'):
-    """ Provide a generator function so that all of the enhanced images don't need
-    to be stored in memory at the same time. This can save tons of RAM compared to
-    the enhancer function. """
-
-    print('face enhancer....')
-    if not isinstance(images, list) and os.path.isfile(images): # handle video to images
-        images = load_video_to_cv2(images)
-
-    # ------------------------ set up GFPGAN restorer ------------------------
-    if  method == 'gfpgan':
-        arch = 'clean'
-        channel_multiplier = 2
-        model_name = 'GFPGANv1.4'
-        url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'
-    elif method == 'RestoreFormer':
-        arch = 'RestoreFormer'
-        channel_multiplier = 2
-        model_name = 'RestoreFormer'
-        url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth'
-    elif method == 'codeformer': # TODO:
-        arch = 'CodeFormer'
-        channel_multiplier = 2
-        model_name = 'CodeFormer'
-        url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
-    else:
-        raise ValueError(f'Wrong model version {method}.')
-
-
-    # ------------------------ set up background upsampler ------------------------
-    if bg_upsampler == 'realesrgan':
-        if not torch.cuda.is_available():  # CPU
-            import warnings
-            warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
-                          'If you really want to use it, please modify the corresponding codes.')
-            bg_upsampler = None
-        else:
-            from basicsr.archs.rrdbnet_arch import RRDBNet
-            from realesrgan import RealESRGANer
-            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-            bg_upsampler = RealESRGANer(
-                scale=2,
-                model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
-                model=model,
-                tile=400,
-                tile_pad=10,
-                pre_pad=0,
-                half=True)  # need to set False in CPU mode
-    else:
-        bg_upsampler = None
-
-    # determine model paths
-    model_path = os.path.join('gfpgan/weights', model_name + '.pth')
-    
-    if not os.path.isfile(model_path):
-        model_path = os.path.join('checkpoints', model_name + '.pth')
-    
-    if not os.path.isfile(model_path):
-        # download pre-trained models from url
-        model_path = url
-
-    restorer = GFPGANer(
-        model_path=model_path,
-        upscale=2,
-        arch=arch,
-        channel_multiplier=channel_multiplier,
-        bg_upsampler=bg_upsampler)
-
-    # ------------------------ restore ------------------------
-    for idx in tqdm(range(len(images)), 'Face Enhancer:'):
-        
-        img = cv2.cvtColor(images[idx], cv2.COLOR_RGB2BGR)
-        
-        # restore faces and background if necessary
-        cropped_faces, restored_faces, r_img = restorer.enhance(
-            img,
-            has_aligned=False,
-            only_center_face=False,
-            paste_back=True)
-        
-        r_img = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB)
-        yield r_img

code/face_sr/videoio.py

@@ -1,41 +0,0 @@
-import shutil
-import uuid
-
-import os
-
-import cv2
-
-def load_video_to_cv2(input_path):
-    video_stream = cv2.VideoCapture(input_path)
-    fps = video_stream.get(cv2.CAP_PROP_FPS)
-    full_frames = [] 
-    while 1:
-        still_reading, frame = video_stream.read()
-        if not still_reading:
-            video_stream.release()
-            break 
-        full_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    return full_frames
-
-def save_video_with_watermark(video, audio, save_path, watermark=False):
-    temp_file = str(uuid.uuid4())+'.mp4'
-    cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
-    os.system(cmd)
-
-    if watermark is False:
-        shutil.move(temp_file, save_path)
-    else:
-        # watermark
-        try:
-            ##### check if stable-diffusion-webui
-            import webui
-            from modules import paths
-            watarmark_path = paths.script_path+"/extensions/SadTalker/docs/sadtalker_logo.png"
-        except:
-            # get the root path of sadtalker.
-            dir_path = os.path.dirname(os.path.realpath(__file__))
-            watarmark_path = dir_path+"/../../docs/sadtalker_logo.png"
-
-        cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -filter_complex "[1]scale=100:-1[wm];[0][wm]overlay=(main_w-overlay_w)-10:10" "%s"' % (temp_file, watarmark_path, save_path)
-        os.system(cmd)
-        os.remove(temp_file)

code/model/__init__.py

@@ -1,6 +0,0 @@
-from typing import Union
-from .unet import BeatGANsUNetModel, BeatGANsUNetConfig
-from .unet_autoenc import BeatGANsAutoencConfig, BeatGANsAutoencModel
-
-Model = Union[BeatGANsUNetModel, BeatGANsAutoencModel]
-ModelConfig = Union[BeatGANsUNetConfig, BeatGANsAutoencConfig]

code/model/base.py

@@ -1,37 +0,0 @@
-# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
-# This program is free software; you can redistribute it and/or modify
-# it under the terms of the MIT License.
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# MIT License for more details.
-
-import numpy as np
-import torch
-
-
-class BaseModule(torch.nn.Module):
-    def __init__(self):
-        super(BaseModule, self).__init__()
-
-    @property
-    def nparams(self):
-        """
-        Returns number of trainable parameters of the module.
-        """
-        num_params = 0
-        for name, param in self.named_parameters():
-            if param.requires_grad:
-                num_params += np.prod(param.detach().cpu().numpy().shape)
-        return num_params
-
-
-    def relocate_input(self, x: list):
-        """
-        Relocates provided tensors to the same device set for the module.
-        """
-        device = next(self.parameters()).device
-        for i in range(len(x)):
-            if isinstance(x[i], torch.Tensor) and x[i].device != device:
-                x[i] = x[i].to(device)
-        return x

code/model/blocks.py

@@ -1,567 +0,0 @@
-import math
-from abc import abstractmethod
-from dataclasses import dataclass
-from numbers import Number
-
-import torch as th
-import torch.nn.functional as F
-from choices import *
-from config_base import BaseConfig
-from torch import nn
-
-from .nn import (avg_pool_nd, conv_nd, linear, normalization,
-                 timestep_embedding, torch_checkpoint, zero_module)
-
-
-class ScaleAt(Enum):
-    after_norm = 'afternorm'
-
-
-class TimestepBlock(nn.Module):
-    """
-    Any module where forward() takes timestep embeddings as a second argument.
-    """
-    @abstractmethod
-    def forward(self, x, emb=None, cond=None, lateral=None):
-        """
-        Apply the module to `x` given `emb` timestep embeddings.
-        """
-
-
-class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
-    """
-    A sequential module that passes timestep embeddings to the children that
-    support it as an extra input.
-    """
-    def forward(self, x, emb=None, cond=None, lateral=None):
-        for layer in self:
-            if isinstance(layer, TimestepBlock):
-                x = layer(x, emb=emb, cond=cond, lateral=lateral)
-            else:
-                x = layer(x)
-        return x
-
-
-@dataclass
-class ResBlockConfig(BaseConfig):
-    channels: int
-    emb_channels: int
-    dropout: float
-    out_channels: int = None
-    # condition the resblock with time (and encoder's output)
-    use_condition: bool = True
-    # whether to use 3x3 conv for skip path when the channels aren't matched
-    use_conv: bool = False
-    # dimension of conv (always 2 = 2d)
-    dims: int = 2
-    # gradient checkpoint
-    use_checkpoint: bool = False
-    up: bool = False
-    down: bool = False
-    # whether to condition with both time & encoder's output
-    two_cond: bool = False
-    # number of encoders' output channels
-    cond_emb_channels: int = None
-    # suggest: False
-    has_lateral: bool = False
-    lateral_channels: int = None
-    # whether to init the convolution with zero weights
-    # this is default from BeatGANs and seems to help learning
-    use_zero_module: bool = True
-
-    def __post_init__(self):
-        self.out_channels = self.out_channels or self.channels
-        self.cond_emb_channels = self.cond_emb_channels or self.emb_channels
-
-    def make_model(self):
-        return ResBlock(self)
-
-
-class ResBlock(TimestepBlock):
-    """
-    A residual block that can optionally change the number of channels.
-
-    total layers:
-        in_layers
-        - norm
-        - act
-        - conv
-        out_layers
-        - norm
-        - (modulation)
-        - act
-        - conv
-    """
-    def __init__(self, conf: ResBlockConfig):
-        super().__init__()
-        self.conf = conf
-
-        #############################
-        # IN LAYERS
-        #############################
-        assert conf.lateral_channels is None
-        layers = [
-            normalization(conf.channels),
-            nn.SiLU(),
-            conv_nd(conf.dims, conf.channels, conf.out_channels, 3, padding=1)
-        ]
-        self.in_layers = nn.Sequential(*layers)
-
-        self.updown = conf.up or conf.down
-
-        if conf.up:
-            self.h_upd = Upsample(conf.channels, False, conf.dims)
-            self.x_upd = Upsample(conf.channels, False, conf.dims)
-        elif conf.down:
-            self.h_upd = Downsample(conf.channels, False, conf.dims)
-            self.x_upd = Downsample(conf.channels, False, conf.dims)
-        else:
-            self.h_upd = self.x_upd = nn.Identity()
-
-        #############################
-        # OUT LAYERS CONDITIONS
-        #############################
-        if conf.use_condition:
-            # condition layers for the out_layers
-            self.emb_layers = nn.Sequential(
-                nn.SiLU(),
-                linear(conf.emb_channels, 2 * conf.out_channels),
-            )
-
-            if conf.two_cond:
-                self.cond_emb_layers = nn.Sequential(
-                    nn.SiLU(),
-                    linear(conf.cond_emb_channels, conf.out_channels),
-                )
-            #############################
-            # OUT LAYERS (ignored when there is no condition)
-            #############################
-            # original version
-            conv = conv_nd(conf.dims,
-                           conf.out_channels,
-                           conf.out_channels,
-                           3,
-                           padding=1)
-            if conf.use_zero_module:
-                # zere out the weights
-                # it seems to help training
-                conv = zero_module(conv)
-
-            # construct the layers
-            # - norm
-            # - (modulation)
-            # - act
-            # - dropout
-            # - conv
-            layers = []
-            layers += [
-                normalization(conf.out_channels),
-                nn.SiLU(),
-                nn.Dropout(p=conf.dropout),
-                conv,
-            ]
-            self.out_layers = nn.Sequential(*layers)
-
-        #############################
-        # SKIP LAYERS
-        #############################
-        if conf.out_channels == conf.channels:
-            # cannot be used with gatedconv, also gatedconv is alsways used as the first block
-            self.skip_connection = nn.Identity()
-        else:
-            if conf.use_conv:
-                kernel_size = 3
-                padding = 1
-            else:
-                kernel_size = 1
-                padding = 0
-
-            self.skip_connection = conv_nd(conf.dims,
-                                           conf.channels,
-                                           conf.out_channels,
-                                           kernel_size,
-                                           padding=padding)
-
-    def forward(self, x, emb=None, cond=None, lateral=None):
-        """
-        Apply the block to a Tensor, conditioned on a timestep embedding.
-
-        Args:
-            x: input
-            lateral: lateral connection from the encoder
-        """
-        return torch_checkpoint(self._forward, (x, emb, cond, lateral),
-                                self.conf.use_checkpoint)
-
-    def _forward(
-        self,
-        x,
-        emb=None,
-        cond=None,
-        lateral=None,
-    ):
-        """
-        Args:
-            lateral: required if "has_lateral" and non-gated, with gated, it can be supplied optionally    
-        """
-        if self.conf.has_lateral:
-            # lateral may be supplied even if it doesn't require
-            # the model will take the lateral only if "has_lateral"
-            assert lateral is not None
-            x = th.cat([x, lateral], dim=1)
-
-        if self.updown:
-            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
-            h = in_rest(x)
-            h = self.h_upd(h)
-            x = self.x_upd(x)
-            h = in_conv(h)
-        else:
-            h = self.in_layers(x)
-
-        if self.conf.use_condition:
-            # it's possible that the network may not receieve the time emb
-            # this happens with autoenc and setting the time_at
-            if emb is not None:
-                emb_out = self.emb_layers(emb).type(h.dtype)
-            else:
-                emb_out = None
-
-            if self.conf.two_cond:
-                # it's possible that the network is two_cond
-                # but it doesn't get the second condition
-                # in which case, we ignore the second condition
-                # and treat as if the network has one condition
-                if cond is None:
-                    cond_out = None
-                else:
-                    cond_out = self.cond_emb_layers(cond).type(h.dtype)
-
-                if cond_out is not None:
-                    while len(cond_out.shape) < len(h.shape):
-                        cond_out = cond_out[..., None]
-            else:
-                cond_out = None
-
-            # this is the new refactored code
-            h = apply_conditions(
-                h=h,
-                emb=emb_out,
-                cond=cond_out,
-                layers=self.out_layers,
-                scale_bias=1,
-                in_channels=self.conf.out_channels,
-                up_down_layer=None,
-            )
-
-        return self.skip_connection(x) + h
-
-
-def apply_conditions(
-    h,
-    emb=None,
-    cond=None,
-    layers: nn.Sequential = None,
-    scale_bias: float = 1,
-    in_channels: int = 512,
-    up_down_layer: nn.Module = None,
-):
-    """
-    apply conditions on the feature maps
-
-    Args:
-        emb: time conditional (ready to scale + shift)
-        cond: encoder's conditional (read to scale + shift)
-    """
-    two_cond = emb is not None and cond is not None
-
-    if emb is not None:
-        # adjusting shapes
-        while len(emb.shape) < len(h.shape):
-            emb = emb[..., None]
-
-    if two_cond:
-        # adjusting shapes
-        while len(cond.shape) < len(h.shape):
-            cond = cond[..., None]
-        # time first
-        scale_shifts = [emb, cond]
-    else:
-        # "cond" is not used with single cond mode
-        scale_shifts = [emb]
-
-    # support scale, shift or shift only
-    for i, each in enumerate(scale_shifts):
-        if each is None:
-            # special case: the condition is not provided
-            a = None
-            b = None
-        else:
-            if each.shape[1] == in_channels * 2:
-                a, b = th.chunk(each, 2, dim=1)
-            else:
-                a = each
-                b = None
-        scale_shifts[i] = (a, b)
-
-    # condition scale bias could be a list
-    if isinstance(scale_bias, Number):
-        biases = [scale_bias] * len(scale_shifts)
-    else:
-        # a list
-        biases = scale_bias
-
-    # default, the scale & shift are applied after the group norm but BEFORE SiLU
-    pre_layers, post_layers = layers[0], layers[1:]
-
-    # spilt the post layer to be able to scale up or down before conv
-    # post layers will contain only the conv
-    mid_layers, post_layers = post_layers[:-2], post_layers[-2:]
-
-    h = pre_layers(h)
-    # scale and shift for each condition
-    for i, (scale, shift) in enumerate(scale_shifts):
-        # if scale is None, it indicates that the condition is not provided
-        if scale is not None:
-            h = h * (biases[i] + scale)
-            if shift is not None:
-                h = h + shift
-    h = mid_layers(h)
-
-    # upscale or downscale if any just before the last conv
-    if up_down_layer is not None:
-        h = up_down_layer(h)
-    h = post_layers(h)
-    return h
-
-
-class Upsample(nn.Module):
-    """
-    An upsampling layer with an optional convolution.
-
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 upsampling occurs in the inner-two dimensions.
-    """
-    def __init__(self, channels, use_conv, dims=2, out_channels=None):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        if use_conv:
-            self.conv = conv_nd(dims,
-                                self.channels,
-                                self.out_channels,
-                                3,
-                                padding=1)
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        if self.dims == 3:
-            x = F.interpolate(x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2),
-                              mode="nearest")
-        else:
-            x = F.interpolate(x, scale_factor=2, mode="nearest")
-        if self.use_conv:
-            x = self.conv(x)
-        return x
-
-
-class Downsample(nn.Module):
-    """
-    A downsampling layer with an optional convolution.
-
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 downsampling occurs in the inner-two dimensions.
-    """
-    def __init__(self, channels, use_conv, dims=2, out_channels=None):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        stride = 2 if dims != 3 else (1, 2, 2)
-        if use_conv:
-            self.op = conv_nd(dims,
-                              self.channels,
-                              self.out_channels,
-                              3,
-                              stride=stride,
-                              padding=1)
-        else:
-            assert self.channels == self.out_channels
-            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        return self.op(x)
-
-
-class AttentionBlock(nn.Module):
-    """
-    An attention block that allows spatial positions to attend to each other.
-
-    Originally ported from here, but adapted to the N-d case.
-    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
-    """
-    def __init__(
-        self,
-        channels,
-        num_heads=1,
-        num_head_channels=-1,
-        use_checkpoint=False,
-        use_new_attention_order=False,
-    ):
-        super().__init__()
-        self.channels = channels
-        if num_head_channels == -1:
-            self.num_heads = num_heads
-        else:
-            assert (
-                channels % num_head_channels == 0
-            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
-            self.num_heads = channels // num_head_channels
-        self.use_checkpoint = use_checkpoint
-        self.norm = normalization(channels)
-        self.qkv = conv_nd(1, channels, channels * 3, 1)
-        if use_new_attention_order:
-            # split qkv before split heads
-            self.attention = QKVAttention(self.num_heads)
-        else:
-            # split heads before split qkv
-            self.attention = QKVAttentionLegacy(self.num_heads)
-
-        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
-
-    def forward(self, x):
-        return torch_checkpoint(self._forward, (x, ), self.use_checkpoint)
-
-    def _forward(self, x):
-        b, c, *spatial = x.shape
-        x = x.reshape(b, c, -1)
-        qkv = self.qkv(self.norm(x))
-        h = self.attention(qkv)
-        h = self.proj_out(h)
-        return (x + h).reshape(b, c, *spatial)
-
-
-def count_flops_attn(model, _x, y):
-    """
-    A counter for the `thop` package to count the operations in an
-    attention operation.
-    Meant to be used like:
-        macs, params = thop.profile(
-            model,
-            inputs=(inputs, timestamps),
-            custom_ops={QKVAttention: QKVAttention.count_flops},
-        )
-    """
-    b, c, *spatial = y[0].shape
-    num_spatial = int(np.prod(spatial))
-    # We perform two matmuls with the same number of ops.
-    # The first computes the weight matrix, the second computes
-    # the combination of the value vectors.
-    matmul_ops = 2 * b * (num_spatial**2) * c
-    model.total_ops += th.DoubleTensor([matmul_ops])
-
-
-class QKVAttentionLegacy(nn.Module):
-    """
-    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
-    """
-    def __init__(self, n_heads):
-        super().__init__()
-        self.n_heads = n_heads
-
-    def forward(self, qkv):
-        """
-        Apply QKV attention.
-
-        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
-        :return: an [N x (H * C) x T] tensor after attention.
-        """
-        bs, width, length = qkv.shape
-        assert width % (3 * self.n_heads) == 0
-        ch = width // (3 * self.n_heads)
-        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch,
-                                                                       dim=1)
-        scale = 1 / math.sqrt(math.sqrt(ch))
-        weight = th.einsum(
-            "bct,bcs->bts", q * scale,
-            k * scale)  # More stable with f16 than dividing afterwards
-        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
-        a = th.einsum("bts,bcs->bct", weight, v)
-        return a.reshape(bs, -1, length)
-
-    @staticmethod
-    def count_flops(model, _x, y):
-        return count_flops_attn(model, _x, y)
-
-
-class QKVAttention(nn.Module):
-    """
-    A module which performs QKV attention and splits in a different order.
-    """
-    def __init__(self, n_heads):
-        super().__init__()
-        self.n_heads = n_heads
-
-    def forward(self, qkv):
-        """
-        Apply QKV attention.
-
-        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
-        :return: an [N x (H * C) x T] tensor after attention.
-        """
-        bs, width, length = qkv.shape
-        assert width % (3 * self.n_heads) == 0
-        ch = width // (3 * self.n_heads)
-        q, k, v = qkv.chunk(3, dim=1)
-        scale = 1 / math.sqrt(math.sqrt(ch))
-        weight = th.einsum(
-            "bct,bcs->bts",
-            (q * scale).view(bs * self.n_heads, ch, length),
-            (k * scale).view(bs * self.n_heads, ch, length),
-        )  # More stable with f16 than dividing afterwards
-        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
-        a = th.einsum("bts,bcs->bct", weight,
-                      v.reshape(bs * self.n_heads, ch, length))
-        return a.reshape(bs, -1, length)
-
-    @staticmethod
-    def count_flops(model, _x, y):
-        return count_flops_attn(model, _x, y)
-
-
-class AttentionPool2d(nn.Module):
-    """
-    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
-    """
-    def __init__(
-        self,
-        spacial_dim: int,
-        embed_dim: int,
-        num_heads_channels: int,
-        output_dim: int = None,
-    ):
-        super().__init__()
-        self.positional_embedding = nn.Parameter(
-            th.randn(embed_dim, spacial_dim**2 + 1) / embed_dim**0.5)
-        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
-        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
-        self.num_heads = embed_dim // num_heads_channels
-        self.attention = QKVAttention(self.num_heads)
-
-    def forward(self, x):
-        b, c, *_spatial = x.shape
-        x = x.reshape(b, c, -1)  # NC(HW)
-        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)
-        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)
-        x = self.qkv_proj(x)
-        x = self.attention(x)
-        x = self.c_proj(x)
-        return x[:, :, 0]

code/model/diffusion.py

@@ -1,294 +0,0 @@
-# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
-# This program is free software; you can redistribute it and/or modify
-# it under the terms of the MIT License.
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# MIT License for more details.
-
-import math
-import torch
-from einops import rearrange
-
-from model.base import BaseModule
-
-
-class Mish(BaseModule):
-    def forward(self, x):
-        return x * torch.tanh(torch.nn.functional.softplus(x))
-
-
-class Upsample(BaseModule):
-    def __init__(self, dim):
-        super(Upsample, self).__init__()
-        self.conv = torch.nn.ConvTranspose2d(dim, dim, 4, 2, 1)
-
-    def forward(self, x):
-        return self.conv(x)
-
-
-class Downsample(BaseModule):
-    def __init__(self, dim):
-        super(Downsample, self).__init__()
-        self.conv = torch.nn.Conv2d(dim, dim, 3, 2, 1)
-
-    def forward(self, x):
-        return self.conv(x)
-
-
-class Rezero(BaseModule):
-    def __init__(self, fn):
-        super(Rezero, self).__init__()
-        self.fn = fn
-        self.g = torch.nn.Parameter(torch.zeros(1))
-
-    def forward(self, x):
-        return self.fn(x) * self.g
-
-
-class Block(BaseModule):
-    def __init__(self, dim, dim_out, groups=8):
-        super(Block, self).__init__()
-        self.block = torch.nn.Sequential(torch.nn.Conv2d(dim, dim_out, 3, 
-                                         padding=1), torch.nn.GroupNorm(
-                                         groups, dim_out), Mish())
-
-    def forward(self, x, mask):
-        output = self.block(x * mask)
-        return output * mask
-
-
-class ResnetBlock(BaseModule):
-    def __init__(self, dim, dim_out, time_emb_dim, groups=8):
-        super(ResnetBlock, self).__init__()
-        self.mlp = torch.nn.Sequential(Mish(), torch.nn.Linear(time_emb_dim, 
-                                                               dim_out))
-
-        self.block1 = Block(dim, dim_out, groups=groups)
-        self.block2 = Block(dim_out, dim_out, groups=groups)
-        if dim != dim_out:
-            self.res_conv = torch.nn.Conv2d(dim, dim_out, 1)
-        else:
-            self.res_conv = torch.nn.Identity()
-
-    def forward(self, x, mask, time_emb):
-        h = self.block1(x, mask)
-        h += self.mlp(time_emb).unsqueeze(-1).unsqueeze(-1)
-        h = self.block2(h, mask)
-        output = h + self.res_conv(x * mask)
-        return output
-
-
-class LinearAttention(BaseModule):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super(LinearAttention, self).__init__()
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = torch.nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_out = torch.nn.Conv2d(hidden_dim, dim, 1)            
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', 
-                            heads = self.heads, qkv=3)            
-        k = k.softmax(dim=-1)
-        context = torch.einsum('bhdn,bhen->bhde', k, v)
-        out = torch.einsum('bhde,bhdn->bhen', context, q)
-        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', 
-                        heads=self.heads, h=h, w=w)
-        return self.to_out(out)
-
-
-class Residual(BaseModule):
-    def __init__(self, fn):
-        super(Residual, self).__init__()
-        self.fn = fn
-
-    def forward(self, x, *args, **kwargs):
-        output = self.fn(x, *args, **kwargs) + x
-        return output
-
-
-class SinusoidalPosEmb(BaseModule):
-    def __init__(self, dim):
-        super(SinusoidalPosEmb, self).__init__()
-        self.dim = dim
-
-    def forward(self, x, scale=1000):
-        device = x.device
-        half_dim = self.dim // 2
-        emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
-        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
-        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
-        return emb
-
-
-class GradLogPEstimator2d(BaseModule):
-    def __init__(self, dim, dim_mults=(1, 2, 4), groups=8,
-                 n_spks=None, spk_emb_dim=64, n_feats=80, pe_scale=1000):
-        super(GradLogPEstimator2d, self).__init__()
-        self.dim = dim
-        self.dim_mults = dim_mults
-        self.groups = groups
-        self.n_spks = n_spks if not isinstance(n_spks, type(None)) else 1
-        self.spk_emb_dim = spk_emb_dim
-        self.pe_scale = pe_scale
-        
-        if n_spks > 1:
-            self.spk_mlp = torch.nn.Sequential(torch.nn.Linear(spk_emb_dim, spk_emb_dim * 4), Mish(),
-                                               torch.nn.Linear(spk_emb_dim * 4, n_feats))
-        self.time_pos_emb = SinusoidalPosEmb(dim)
-        self.mlp = torch.nn.Sequential(torch.nn.Linear(dim, dim * 4), Mish(),
-                                       torch.nn.Linear(dim * 4, dim))
-
-        dims = [2 + (1 if n_spks > 1 else 0), *map(lambda m: dim * m, dim_mults)]
-        in_out = list(zip(dims[:-1], dims[1:]))
-        self.downs = torch.nn.ModuleList([])
-        self.ups = torch.nn.ModuleList([])
-        num_resolutions = len(in_out)
-
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-            self.downs.append(torch.nn.ModuleList([
-                       ResnetBlock(dim_in, dim_out, time_emb_dim=dim),
-                       ResnetBlock(dim_out, dim_out, time_emb_dim=dim),
-                       Residual(Rezero(LinearAttention(dim_out))),
-                       Downsample(dim_out) if not is_last else torch.nn.Identity()]))
-
-        mid_dim = dims[-1]
-        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)
-        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
-        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)
-
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            self.ups.append(torch.nn.ModuleList([
-                     ResnetBlock(dim_out * 2, dim_in, time_emb_dim=dim),
-                     ResnetBlock(dim_in, dim_in, time_emb_dim=dim),
-                     Residual(Rezero(LinearAttention(dim_in))),
-                     Upsample(dim_in)]))
-        self.final_block = Block(dim, dim)
-        self.final_conv = torch.nn.Conv2d(dim, 1, 1)
-
-    def forward(self, x, mask, mu, t, spk=None):
-        if not isinstance(spk, type(None)):
-            s = self.spk_mlp(spk)
-        
-        t = self.time_pos_emb(t, scale=self.pe_scale)
-        t = self.mlp(t)
-
-        if self.n_spks < 2:
-            x = torch.stack([mu, x], 1)
-        else:
-            s = s.unsqueeze(-1).repeat(1, 1, x.shape[-1])
-            x = torch.stack([mu, x, s], 1)
-        mask = mask.unsqueeze(1)
-
-        hiddens = []
-        masks = [mask]
-        for resnet1, resnet2, attn, downsample in self.downs:
-            mask_down = masks[-1]
-            x = resnet1(x, mask_down, t)
-            x = resnet2(x, mask_down, t)
-            x = attn(x)
-            hiddens.append(x)
-            x = downsample(x * mask_down)
-            masks.append(mask_down[:, :, :, ::2])
-
-        masks = masks[:-1]
-        mask_mid = masks[-1]
-        x = self.mid_block1(x, mask_mid, t)
-        x = self.mid_attn(x)
-        x = self.mid_block2(x, mask_mid, t)
-
-        for resnet1, resnet2, attn, upsample in self.ups:
-            mask_up = masks.pop()
-            x = torch.cat((x, hiddens.pop()), dim=1)
-            x = resnet1(x, mask_up, t)
-            x = resnet2(x, mask_up, t)
-            x = attn(x)
-            x = upsample(x * mask_up)
-
-        x = self.final_block(x, mask)
-        output = self.final_conv(x * mask)
-
-        return (output * mask).squeeze(1)
-
-
-def get_noise(t, beta_init, beta_term, cumulative=False):
-    if cumulative:
-        noise = beta_init*t + 0.5*(beta_term - beta_init)*(t**2)
-    else:
-        noise = beta_init + (beta_term - beta_init)*t
-    return noise
-
-
-class Diffusion(BaseModule):
-    def __init__(self, n_feats, dim,
-                 n_spks=1, spk_emb_dim=64,
-                 beta_min=0.05, beta_max=20, pe_scale=1000):
-        super(Diffusion, self).__init__()
-        self.n_feats = n_feats
-        self.dim = dim
-        self.n_spks = n_spks
-        self.spk_emb_dim = spk_emb_dim
-        self.beta_min = beta_min
-        self.beta_max = beta_max
-        self.pe_scale = pe_scale
-        
-        self.estimator = GradLogPEstimator2d(dim, n_spks=n_spks,
-                                             spk_emb_dim=spk_emb_dim,
-                                             pe_scale=pe_scale)
-
-    def forward_diffusion(self, x0, mask, mu, t):
-        time = t.unsqueeze(-1).unsqueeze(-1)
-        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True)
-        mean = x0*torch.exp(-0.5*cum_noise) + mu*(1.0 - torch.exp(-0.5*cum_noise))
-        variance = 1.0 - torch.exp(-cum_noise)
-        z = torch.randn(x0.shape, dtype=x0.dtype, device=x0.device, 
-                        requires_grad=False)
-        xt = mean + z * torch.sqrt(variance)
-        return xt * mask, z * mask
-
-    @torch.no_grad()
-    def reverse_diffusion(self, z, mask, mu, n_timesteps, stoc=False, spk=None):
-        h = 1.0 / n_timesteps
-        xt = z * mask
-        for i in range(n_timesteps):
-            t = (1.0 - (i + 0.5)*h) * torch.ones(z.shape[0], dtype=z.dtype, 
-                                                 device=z.device)
-            time = t.unsqueeze(-1).unsqueeze(-1)
-            noise_t = get_noise(time, self.beta_min, self.beta_max, 
-                                cumulative=False)
-            if stoc:  # adds stochastic term
-                dxt_det = 0.5 * (mu - xt) - self.estimator(xt, mask, mu, t, spk)
-                dxt_det = dxt_det * noise_t * h
-                dxt_stoc = torch.randn(z.shape, dtype=z.dtype, device=z.device,
-                                       requires_grad=False)
-                dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h)
-                dxt = dxt_det + dxt_stoc
-            else:
-                dxt = 0.5 * (mu - xt - self.estimator(xt, mask, mu, t, spk))
-                dxt = dxt * noise_t * h
-            xt = (xt - dxt) * mask
-        return xt
-
-    @torch.no_grad()
-    def forward(self, z, mask, mu, n_timesteps, stoc=False, spk=None):
-        return self.reverse_diffusion(z, mask, mu, n_timesteps, stoc, spk)
-
-    def loss_t(self, x0, mask, mu, t, spk=None):
-        xt, z = self.forward_diffusion(x0, mask, mu, t)
-        time = t.unsqueeze(-1).unsqueeze(-1)
-        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True)
-        noise_estimation = self.estimator(xt, mask, mu, t, spk)
-        noise_estimation *= torch.sqrt(1.0 - torch.exp(-cum_noise))
-        loss = torch.sum((noise_estimation + z)**2) / (torch.sum(mask)*self.n_feats)
-        return loss, xt
-
-    def compute_loss(self, x0, mask, mu, spk=None, offset=1e-5):
-        t = torch.rand(x0.shape[0], dtype=x0.dtype, device=x0.device,
-                       requires_grad=False)
-        t = torch.clamp(t, offset, 1.0 - offset)
-        return self.loss_t(x0, mask, mu, t, spk)

code/model/latentnet.py

@@ -1,193 +0,0 @@
-import math
-from dataclasses import dataclass
-from enum import Enum
-from typing import NamedTuple, Tuple
-
-import torch
-from choices import *
-from config_base import BaseConfig
-from torch import nn
-from torch.nn import init
-
-from .blocks import *
-from .nn import timestep_embedding
-from .unet import *
-
-
-class LatentNetType(Enum):
-    none = 'none'
-    # injecting inputs into the hidden layers
-    skip = 'skip'
-
-
-class LatentNetReturn(NamedTuple):
-    pred: torch.Tensor = None
-
-
-@dataclass
-class MLPSkipNetConfig(BaseConfig):
-    """
-    default MLP for the latent DPM in the paper!
-    """
-    num_channels: int
-    skip_layers: Tuple[int]
-    num_hid_channels: int
-    num_layers: int
-    num_time_emb_channels: int = 64
-    activation: Activation = Activation.silu
-    use_norm: bool = True
-    condition_bias: float = 1
-    dropout: float = 0
-    last_act: Activation = Activation.none
-    num_time_layers: int = 2
-    time_last_act: bool = False
-
-    def make_model(self):
-        return MLPSkipNet(self)
-
-
-class MLPSkipNet(nn.Module):
-    """
-    concat x to hidden layers
-
-    default MLP for the latent DPM in the paper!
-    """
-    def __init__(self, conf: MLPSkipNetConfig):
-        super().__init__()
-        self.conf = conf
-
-        layers = []
-        for i in range(conf.num_time_layers):
-            if i == 0:
-                a = conf.num_time_emb_channels
-                b = conf.num_channels
-            else:
-                a = conf.num_channels
-                b = conf.num_channels
-            layers.append(nn.Linear(a, b))
-            if i < conf.num_time_layers - 1 or conf.time_last_act:
-                layers.append(conf.activation.get_act())
-        self.time_embed = nn.Sequential(*layers)
-
-        self.layers = nn.ModuleList([])
-        for i in range(conf.num_layers):
-            if i == 0:
-                act = conf.activation
-                norm = conf.use_norm
-                cond = True
-                a, b = conf.num_channels, conf.num_hid_channels
-                dropout = conf.dropout
-            elif i == conf.num_layers - 1:
-                act = Activation.none
-                norm = False
-                cond = False
-                a, b = conf.num_hid_channels, conf.num_channels
-                dropout = 0
-            else:
-                act = conf.activation
-                norm = conf.use_norm
-                cond = True
-                a, b = conf.num_hid_channels, conf.num_hid_channels
-                dropout = conf.dropout
-
-            if i in conf.skip_layers:
-                a += conf.num_channels
-
-            self.layers.append(
-                MLPLNAct(
-                    a,
-                    b,
-                    norm=norm,
-                    activation=act,
-                    cond_channels=conf.num_channels,
-                    use_cond=cond,
-                    condition_bias=conf.condition_bias,
-                    dropout=dropout,
-                ))
-        self.last_act = conf.last_act.get_act()
-
-    def forward(self, x, t, **kwargs):
-        t = timestep_embedding(t, self.conf.num_time_emb_channels)
-        cond = self.time_embed(t)
-        h = x
-        for i in range(len(self.layers)):
-            if i in self.conf.skip_layers:
-                # injecting input into the hidden layers
-                h = torch.cat([h, x], dim=1)
-            h = self.layers[i].forward(x=h, cond=cond)
-        h = self.last_act(h)
-        return LatentNetReturn(h)
-
-
-class MLPLNAct(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        norm: bool,
-        use_cond: bool,
-        activation: Activation,
-        cond_channels: int,
-        condition_bias: float = 0,
-        dropout: float = 0,
-    ):
-        super().__init__()
-        self.activation = activation
-        self.condition_bias = condition_bias
-        self.use_cond = use_cond
-
-        self.linear = nn.Linear(in_channels, out_channels)
-        self.act = activation.get_act()
-        if self.use_cond:
-            self.linear_emb = nn.Linear(cond_channels, out_channels)
-            self.cond_layers = nn.Sequential(self.act, self.linear_emb)
-        if norm:
-            self.norm = nn.LayerNorm(out_channels)
-        else:
-            self.norm = nn.Identity()
-
-        if dropout > 0:
-            self.dropout = nn.Dropout(p=dropout)
-        else:
-            self.dropout = nn.Identity()
-
-        self.init_weights()
-
-    def init_weights(self):
-        for module in self.modules():
-            if isinstance(module, nn.Linear):
-                if self.activation == Activation.relu:
-                    init.kaiming_normal_(module.weight,
-                                         a=0,
-                                         nonlinearity='relu')
-                elif self.activation == Activation.lrelu:
-                    init.kaiming_normal_(module.weight,
-                                         a=0.2,
-                                         nonlinearity='leaky_relu')
-                elif self.activation == Activation.silu:
-                    init.kaiming_normal_(module.weight,
-                                         a=0,
-                                         nonlinearity='relu')
-                else:
-                    # leave it as default
-                    pass
-
-    def forward(self, x, cond=None):
-        x = self.linear(x)
-        if self.use_cond:
-            # (n, c) or (n, c * 2)
-            cond = self.cond_layers(cond)
-            cond = (cond, None)
-
-            # scale shift first
-            x = x * (self.condition_bias + cond[0])
-            if cond[1] is not None:
-                x = x + cond[1]
-            # then norm
-            x = self.norm(x)
-        else:
-            # no condition
-            x = self.norm(x)
-        x = self.act(x)
-        x = self.dropout(x)
-        return x

code/model/nn.py

@@ -1,137 +0,0 @@
-"""
-Various utilities for neural networks.
-"""
-
-from enum import Enum
-import math
-from typing import Optional
-
-import torch as th
-import torch.nn as nn
-import torch.utils.checkpoint
-
-import torch.nn.functional as F
-
-
-# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
-class SiLU(nn.Module):
-    # @th.jit.script
-    def forward(self, x):
-        return x * th.sigmoid(x)
-
-
-class GroupNorm32(nn.GroupNorm):
-    def forward(self, x):
-        return super().forward(x.float()).type(x.dtype)
-
-
-def conv_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D convolution module.
-    """
-    if dims == 1:
-        return nn.Conv1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.Conv2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.Conv3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-def linear(*args, **kwargs):
-    """
-    Create a linear module.
-    """
-    return nn.Linear(*args, **kwargs)
-
-
-def avg_pool_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D average pooling module.
-    """
-    if dims == 1:
-        return nn.AvgPool1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.AvgPool2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.AvgPool3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-def update_ema(target_params, source_params, rate=0.99):
-    """
-    Update target parameters to be closer to those of source parameters using
-    an exponential moving average.
-
-    :param target_params: the target parameter sequence.
-    :param source_params: the source parameter sequence.
-    :param rate: the EMA rate (closer to 1 means slower).
-    """
-    for targ, src in zip(target_params, source_params):
-        targ.detach().mul_(rate).add_(src, alpha=1 - rate)
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def scale_module(module, scale):
-    """
-    Scale the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().mul_(scale)
-    return module
-
-
-def mean_flat(tensor):
-    """
-    Take the mean over all non-batch dimensions.
-    """
-    return tensor.mean(dim=list(range(1, len(tensor.shape))))
-
-
-def normalization(channels):
-    """
-    Make a standard normalization layer.
-
-    :param channels: number of input channels.
-    :return: an nn.Module for normalization.
-    """
-    return GroupNorm32(min(32, channels), channels)
-
-
-def timestep_embedding(timesteps, dim, max_period=10000):
-    """
-    Create sinusoidal timestep embeddings.
-
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    half = dim // 2
-    freqs = th.exp(-math.log(max_period) *
-                   th.arange(start=0, end=half, dtype=th.float32) /
-                   half).to(device=timesteps.device)
-    args = timesteps[:, None].float() * freqs[None]
-    embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = th.cat(
-            [embedding, th.zeros_like(embedding[:, :1])], dim=-1)
-    return embedding
-
-
-def torch_checkpoint(func, args, flag, preserve_rng_state=False):
-    # torch's gradient checkpoint works with automatic mixed precision, given torch >= 1.8
-    if flag:
-        return torch.utils.checkpoint.checkpoint(
-            func, *args, preserve_rng_state=preserve_rng_state)
-    else:
-        return func(*args)

code/model/seq2seq.py

@@ -1,141 +0,0 @@
-import torch
-from torch import nn
-from model.base import BaseModule
-from espnet.nets.pytorch_backend.conformer.encoder import Encoder as ConformerEncoder
-import torch.nn.functional as F
-
-class LSTM(nn.Module):
-    def __init__(self, motion_dim, output_dim, num_layers=2, hidden_dim=128):
-        super().__init__()
-        self.lstm = nn.LSTM(input_size=motion_dim, hidden_size=hidden_dim,
-                            num_layers=num_layers, batch_first=True)
-        self.fc = nn.Linear(hidden_dim, output_dim)
-
-    def forward(self, x):
-        x, _ = self.lstm(x)
-        return self.fc(x)
-
-class DiffusionPredictor(BaseModule):
-    def __init__(self, conf):
-        super(DiffusionPredictor, self).__init__()
-        
-        self.infer_type = conf.infer_type
-        
-        self.initialize_layers(conf)
-        print(f'infer_type: {self.infer_type}')
-        
-    def create_conformer_encoder(self, attention_dim, num_blocks):
-        return ConformerEncoder(
-            idim=0, attention_dim=attention_dim, attention_heads=2, linear_units=attention_dim,
-            num_blocks=num_blocks, input_layer=None, dropout_rate=0.2, positional_dropout_rate=0.2,
-            attention_dropout_rate=0.2, normalize_before=False, concat_after=False,
-            positionwise_layer_type="linear", positionwise_conv_kernel_size=3, macaron_style=True,
-            pos_enc_layer_type="rel_pos", selfattention_layer_type="rel_selfattn", use_cnn_module=True,
-            cnn_module_kernel=13)
-
-    def initialize_layers(self,  conf, mfcc_dim=39, hubert_dim=1024, speech_layers=4, speech_dim=512, decoder_dim=1024, motion_start_dim=512, HAL_layers=25):
-
-        self.conf = conf
-        # Speech downsampling
-        if self.infer_type.startswith('mfcc'):
-            # from 100 hz to 25 hz
-            self.down_sample1 = nn.Conv1d(mfcc_dim, 256, kernel_size=3, stride=2, padding=1)
-            self.down_sample2 = nn.Conv1d(256, speech_dim, kernel_size=3, stride=2, padding=1)
-        elif self.infer_type.startswith('hubert'):
-            # from 50 hz to 25 hz
-            self.down_sample1 = nn.Conv1d(hubert_dim, speech_dim, kernel_size=3, stride=2, padding=1)
-            
-            self.weights = nn.Parameter(torch.zeros(HAL_layers))
-            self.speech_encoder = self.create_conformer_encoder(speech_dim, speech_layers)
-        else:
-            print('infer_type not supported')
-            
-        # Encoders & Deocoders
-        self.coarse_decoder = self.create_conformer_encoder(decoder_dim, conf.decoder_layers)
-
-        # LSTM predictors for Variance Adapter
-        if self.infer_type != 'hubert_audio_only':
-            self.pose_predictor = LSTM(speech_dim, 3)
-            self.pose_encoder = LSTM(3, speech_dim)
-
-        if 'full_control' in self.infer_type:
-            self.location_predictor = LSTM(speech_dim, 1)
-            self.location_encoder = LSTM(1, speech_dim)
-            self.face_scale_predictor = LSTM(speech_dim, 1)
-            self.face_scale_encoder = LSTM(1, speech_dim)
-        
-        # Linear transformations
-        self.init_code_proj = nn.Sequential(nn.Linear(motion_start_dim, 128))
-        self.noisy_encoder = nn.Sequential(nn.Linear(conf.motion_dim, 128))
-        self.t_encoder = nn.Sequential(nn.Linear(1, 128))
-        self.encoder_direction_code = nn.Linear(conf.motion_dim, 128)
-        
-        self.out_proj = nn.Linear(decoder_dim, conf.motion_dim)
-
-    
-    def forward(self, initial_code, direction_code, seq_input_vector, face_location, face_scale, yaw_pitch_roll, noisy_x, t_emb, control_flag=False):
-        
-        if self.infer_type.startswith('mfcc'):
-            x = self.mfcc_speech_downsample(seq_input_vector)
-        elif self.infer_type.startswith('hubert'):
-            norm_weights = F.softmax(self.weights, dim=-1)
-            weighted_feature = (norm_weights.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) * seq_input_vector).sum(dim=1)
-            x = self.down_sample1(weighted_feature.transpose(1,2)).transpose(1,2)
-            x, _ = self.speech_encoder(x, masks=None)
-        predicted_location, predicted_scale, predicted_pose = face_location, face_scale, yaw_pitch_roll
-        if self.infer_type != 'hubert_audio_only':
-            print(f'pose controllable. control_flag: {control_flag}')
-            x, predicted_location, predicted_scale, predicted_pose = self.adjust_features(x, face_location, face_scale, yaw_pitch_roll, control_flag)
-        concatenated_features = self.combine_features(x, initial_code, direction_code, noisy_x, t_emb) # initial_code and direction_code serve as a motion guide extracted from the reference image. This aims to tell the model what the starting motion should be.
-        outputs = self.decode_features(concatenated_features)
-        return outputs, predicted_location, predicted_scale, predicted_pose
-
-    def mfcc_speech_downsample(self, seq_input_vector):
-        x = self.down_sample1(seq_input_vector.transpose(1,2))
-        return self.down_sample2(x).transpose(1,2)
-
-    def adjust_features(self, x, face_location, face_scale, yaw_pitch_roll, control_flag):
-        predicted_location,  predicted_scale = 0, 0
-        if 'full_control' in self.infer_type:
-            print(f'full controllable. control_flag: {control_flag}')
-            x_residual, predicted_location = self.adjust_location(x, face_location, control_flag)
-            x = x + x_residual
-
-            x_residual, predicted_scale = self.adjust_scale(x, face_scale, control_flag)
-            x = x + x_residual
-
-        x_residual, predicted_pose= self.adjust_pose(x, yaw_pitch_roll, control_flag)
-        x = x + x_residual
-        return x, predicted_location, predicted_scale, predicted_pose
-
-    def adjust_location(self, x, face_location, control_flag):
-        if control_flag:
-            predicted_location = face_location
-        else:
-            predicted_location = self.location_predictor(x)
-        return self.location_encoder(predicted_location), predicted_location
-
-    def adjust_scale(self, x, face_scale, control_flag):
-        if control_flag:
-            predicted_face_scale = face_scale
-        else:
-            predicted_face_scale = self.face_scale_predictor(x)
-        return self.face_scale_encoder(predicted_face_scale), predicted_face_scale
-
-    def adjust_pose(self, x, yaw_pitch_roll, control_flag):
-        if control_flag:
-            predicted_pose = yaw_pitch_roll
-        else:
-            predicted_pose = self.pose_predictor(x)
-        return self.pose_encoder(predicted_pose), predicted_pose
-
-    def combine_features(self, x, initial_code, direction_code, noisy_x, t_emb):
-        init_code_proj = self.init_code_proj(initial_code).unsqueeze(1).repeat(1, x.size(1), 1)
-        noisy_feature = self.noisy_encoder(noisy_x)
-        t_emb_feature = self.t_encoder(t_emb.unsqueeze(1).float()).unsqueeze(1).repeat(1, x.size(1), 1)
-        direction_code_feature = self.encoder_direction_code(direction_code).unsqueeze(1).repeat(1, x.size(1), 1)
-        return torch.cat((x, direction_code_feature, init_code_proj, noisy_feature, t_emb_feature), dim=-1)
-
-    def decode_features(self, concatenated_features):
-        outputs, _ = self.coarse_decoder(concatenated_features, masks=None)
-        return self.out_proj(outputs)

code/model/unet.py

@@ -1,552 +0,0 @@
-import math
-from dataclasses import dataclass
-from numbers import Number
-from typing import NamedTuple, Tuple, Union
-
-import numpy as np
-import torch as th
-from torch import nn
-import torch.nn.functional as F
-from choices import *
-from config_base import BaseConfig
-from .blocks import *
-
-from .nn import (conv_nd, linear, normalization, timestep_embedding,
-                 torch_checkpoint, zero_module)
-
-
-@dataclass
-class BeatGANsUNetConfig(BaseConfig):
-    image_size: int = 64
-    in_channels: int = 3
-    # base channels, will be multiplied
-    model_channels: int = 64
-    # output of the unet
-    # suggest: 3
-    # you only need 6 if you also model the variance of the noise prediction (usually we use an analytical variance hence 3)
-    out_channels: int = 3
-    # how many repeating resblocks per resolution
-    # the decoding side would have "one more" resblock
-    # default: 2
-    num_res_blocks: int = 2
-    # you can also set the number of resblocks specifically for the input blocks
-    # default: None = above
-    num_input_res_blocks: int = None
-    # number of time embed channels and style channels
-    embed_channels: int = 512
-    # at what resolutions you want to do self-attention of the feature maps
-    # attentions generally improve performance
-    # default: [16]
-    # beatgans: [32, 16, 8]
-    attention_resolutions: Tuple[int] = (16, )
-    # number of time embed channels
-    time_embed_channels: int = None
-    # dropout applies to the resblocks (on feature maps)
-    dropout: float = 0.1
-    channel_mult: Tuple[int] = (1, 2, 4, 8)
-    input_channel_mult: Tuple[int] = None
-    conv_resample: bool = True
-    # always 2 = 2d conv
-    dims: int = 2
-    # don't use this, legacy from BeatGANs
-    num_classes: int = None
-    use_checkpoint: bool = False
-    # number of attention heads
-    num_heads: int = 1
-    # or specify the number of channels per attention head
-    num_head_channels: int = -1
-    # what's this?
-    num_heads_upsample: int = -1
-    # use resblock for upscale/downscale blocks (expensive)
-    # default: True (BeatGANs)
-    resblock_updown: bool = True
-    # never tried
-    use_new_attention_order: bool = False
-    resnet_two_cond: bool = False
-    resnet_cond_channels: int = None
-    # init the decoding conv layers with zero weights, this speeds up training
-    # default: True (BeattGANs)
-    resnet_use_zero_module: bool = True
-    # gradient checkpoint the attention operation
-    attn_checkpoint: bool = False
-
-    def make_model(self):
-        return BeatGANsUNetModel(self)
-
-
-class BeatGANsUNetModel(nn.Module):
-    def __init__(self, conf: BeatGANsUNetConfig):
-        super().__init__()
-        self.conf = conf
-
-        if conf.num_heads_upsample == -1:
-            self.num_heads_upsample = conf.num_heads
-
-        self.dtype = th.float32
-
-        self.time_emb_channels = conf.time_embed_channels or conf.model_channels
-        self.time_embed = nn.Sequential(
-            linear(self.time_emb_channels, conf.embed_channels),
-            nn.SiLU(),
-            linear(conf.embed_channels, conf.embed_channels),
-        )
-
-        if conf.num_classes is not None:
-            self.label_emb = nn.Embedding(conf.num_classes,
-                                          conf.embed_channels)
-
-        ch = input_ch = int(conf.channel_mult[0] * conf.model_channels)
-        self.input_blocks = nn.ModuleList([
-            TimestepEmbedSequential(
-                conv_nd(conf.dims, conf.in_channels, ch, 3, padding=1))
-        ])
-
-        kwargs = dict(
-            use_condition=True,
-            two_cond=conf.resnet_two_cond,
-            use_zero_module=conf.resnet_use_zero_module,
-            # style channels for the resnet block
-            cond_emb_channels=conf.resnet_cond_channels,
-        )
-
-        self._feature_size = ch
-
-        # input_block_chans = [ch]
-        input_block_chans = [[] for _ in range(len(conf.channel_mult))]
-        input_block_chans[0].append(ch)
-
-        # number of blocks at each resolution
-        self.input_num_blocks = [0 for _ in range(len(conf.channel_mult))]
-        self.input_num_blocks[0] = 1
-        self.output_num_blocks = [0 for _ in range(len(conf.channel_mult))]
-
-        ds = 1
-        resolution = conf.image_size
-        for level, mult in enumerate(conf.input_channel_mult
-                                     or conf.channel_mult):
-            for _ in range(conf.num_input_res_blocks or conf.num_res_blocks):
-                layers = [
-                    ResBlockConfig(
-                        ch,
-                        conf.embed_channels,
-                        conf.dropout,
-                        out_channels=int(mult * conf.model_channels),
-                        dims=conf.dims,
-                        use_checkpoint=conf.use_checkpoint,
-                        **kwargs,
-                    ).make_model()
-                ]
-                ch = int(mult * conf.model_channels)
-                if resolution in conf.attention_resolutions:
-                    layers.append(
-                        AttentionBlock(
-                            ch,
-                            use_checkpoint=conf.use_checkpoint
-                            or conf.attn_checkpoint,
-                            num_heads=conf.num_heads,
-                            num_head_channels=conf.num_head_channels,
-                            use_new_attention_order=conf.
-                            use_new_attention_order,
-                        ))
-                self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
-                # input_block_chans.append(ch)
-                input_block_chans[level].append(ch)
-                self.input_num_blocks[level] += 1
-                # print(input_block_chans)
-            if level != len(conf.channel_mult) - 1:
-                resolution //= 2
-                out_ch = ch
-                self.input_blocks.append(
-                    TimestepEmbedSequential(
-                        ResBlockConfig(
-                            ch,
-                            conf.embed_channels,
-                            conf.dropout,
-                            out_channels=out_ch,
-                            dims=conf.dims,
-                            use_checkpoint=conf.use_checkpoint,
-                            down=True,
-                            **kwargs,
-                        ).make_model() if conf.
-                        resblock_updown else Downsample(ch,
-                                                        conf.conv_resample,
-                                                        dims=conf.dims,
-                                                        out_channels=out_ch)))
-                ch = out_ch
-                # input_block_chans.append(ch)
-                input_block_chans[level + 1].append(ch)
-                self.input_num_blocks[level + 1] += 1
-                ds *= 2
-                self._feature_size += ch
-
-        self.middle_block = TimestepEmbedSequential(
-            ResBlockConfig(
-                ch,
-                conf.embed_channels,
-                conf.dropout,
-                dims=conf.dims,
-                use_checkpoint=conf.use_checkpoint,
-                **kwargs,
-            ).make_model(),
-            AttentionBlock(
-                ch,
-                use_checkpoint=conf.use_checkpoint or conf.attn_checkpoint,
-                num_heads=conf.num_heads,
-                num_head_channels=conf.num_head_channels,
-                use_new_attention_order=conf.use_new_attention_order,
-            ),
-            ResBlockConfig(
-                ch,
-                conf.embed_channels,
-                conf.dropout,
-                dims=conf.dims,
-                use_checkpoint=conf.use_checkpoint,
-                **kwargs,
-            ).make_model(),
-        )
-        self._feature_size += ch
-
-        self.output_blocks = nn.ModuleList([])
-        for level, mult in list(enumerate(conf.channel_mult))[::-1]:
-            for i in range(conf.num_res_blocks + 1):
-                # print(input_block_chans)
-                # ich = input_block_chans.pop()
-                try:
-                    ich = input_block_chans[level].pop()
-                except IndexError:
-                    # this happens only when num_res_block > num_enc_res_block
-                    # we will not have enough lateral (skip) connecions for all decoder blocks
-                    ich = 0
-                # print('pop:', ich)
-                layers = [
-                    ResBlockConfig(
-                        # only direct channels when gated
-                        channels=ch + ich,
-                        emb_channels=conf.embed_channels,
-                        dropout=conf.dropout,
-                        out_channels=int(conf.model_channels * mult),
-                        dims=conf.dims,
-                        use_checkpoint=conf.use_checkpoint,
-                        # lateral channels are described here when gated
-                        has_lateral=True if ich > 0 else False,
-                        lateral_channels=None,
-                        **kwargs,
-                    ).make_model()
-                ]
-                ch = int(conf.model_channels * mult)
-                if resolution in conf.attention_resolutions:
-                    layers.append(
-                        AttentionBlock(
-                            ch,
-                            use_checkpoint=conf.use_checkpoint
-                            or conf.attn_checkpoint,
-                            num_heads=self.num_heads_upsample,
-                            num_head_channels=conf.num_head_channels,
-                            use_new_attention_order=conf.
-                            use_new_attention_order,
-                        ))
-                if level and i == conf.num_res_blocks:
-                    resolution *= 2
-                    out_ch = ch
-                    layers.append(
-                        ResBlockConfig(
-                            ch,
-                            conf.embed_channels,
-                            conf.dropout,
-                            out_channels=out_ch,
-                            dims=conf.dims,
-                            use_checkpoint=conf.use_checkpoint,
-                            up=True,
-                            **kwargs,
-                        ).make_model() if (
-                            conf.resblock_updown
-                        ) else Upsample(ch,
-                                        conf.conv_resample,
-                                        dims=conf.dims,
-                                        out_channels=out_ch))
-                    ds //= 2
-                self.output_blocks.append(TimestepEmbedSequential(*layers))
-                self.output_num_blocks[level] += 1
-                self._feature_size += ch
-
-        # print(input_block_chans)
-        # print('inputs:', self.input_num_blocks)
-        # print('outputs:', self.output_num_blocks)
-
-        if conf.resnet_use_zero_module:
-            self.out = nn.Sequential(
-                normalization(ch),
-                nn.SiLU(),
-                zero_module(
-                    conv_nd(conf.dims,
-                            input_ch,
-                            conf.out_channels,
-                            3,
-                            padding=1)),
-            )
-        else:
-            self.out = nn.Sequential(
-                normalization(ch),
-                nn.SiLU(),
-                conv_nd(conf.dims, input_ch, conf.out_channels, 3, padding=1),
-            )
-
-    def forward(self, x, t, y=None, **kwargs):
-        """
-        Apply the model to an input batch.
-
-        :param x: an [N x C x ...] Tensor of inputs.
-        :param timesteps: a 1-D batch of timesteps.
-        :param y: an [N] Tensor of labels, if class-conditional.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        assert (y is not None) == (
-            self.conf.num_classes is not None
-        ), "must specify y if and only if the model is class-conditional"
-
-        # hs = []
-        hs = [[] for _ in range(len(self.conf.channel_mult))]
-        emb = self.time_embed(timestep_embedding(t, self.time_emb_channels))
-
-        if self.conf.num_classes is not None:
-            raise NotImplementedError()
-            # assert y.shape == (x.shape[0], )
-            # emb = emb + self.label_emb(y)
-
-        # new code supports input_num_blocks != output_num_blocks
-        h = x.type(self.dtype)
-        k = 0
-        for i in range(len(self.input_num_blocks)):
-            for j in range(self.input_num_blocks[i]):
-                h = self.input_blocks[k](h, emb=emb)
-                # print(i, j, h.shape)
-                hs[i].append(h)
-                k += 1
-        assert k == len(self.input_blocks)
-
-        h = self.middle_block(h, emb=emb)
-        k = 0
-        for i in range(len(self.output_num_blocks)):
-            for j in range(self.output_num_blocks[i]):
-                # take the lateral connection from the same layer (in reserve)
-                # until there is no more, use None
-                try:
-                    lateral = hs[-i - 1].pop()
-                    # print(i, j, lateral.shape)
-                except IndexError:
-                    lateral = None
-                    # print(i, j, lateral)
-                h = self.output_blocks[k](h, emb=emb, lateral=lateral)
-                k += 1
-
-        h = h.type(x.dtype)
-        pred = self.out(h)
-        return Return(pred=pred)
-
-
-class Return(NamedTuple):
-    pred: th.Tensor
-
-
-@dataclass
-class BeatGANsEncoderConfig(BaseConfig):
-    image_size: int
-    in_channels: int
-    model_channels: int
-    out_hid_channels: int
-    out_channels: int
-    num_res_blocks: int
-    attention_resolutions: Tuple[int]
-    dropout: float = 0
-    channel_mult: Tuple[int] = (1, 2, 4, 8)
-    use_time_condition: bool = True
-    conv_resample: bool = True
-    dims: int = 2
-    use_checkpoint: bool = False
-    num_heads: int = 1
-    num_head_channels: int = -1
-    resblock_updown: bool = False
-    use_new_attention_order: bool = False
-    pool: str = 'adaptivenonzero'
-
-    def make_model(self):
-        return BeatGANsEncoderModel(self)
-
-
-class BeatGANsEncoderModel(nn.Module):
-    """
-    The half UNet model with attention and timestep embedding.
-
-    For usage, see UNet.
-    """
-    def __init__(self, conf: BeatGANsEncoderConfig):
-        super().__init__()
-        self.conf = conf
-        self.dtype = th.float32
-
-        if conf.use_time_condition:
-            time_embed_dim = conf.model_channels * 4
-            self.time_embed = nn.Sequential(
-                linear(conf.model_channels, time_embed_dim),
-                nn.SiLU(),
-                linear(time_embed_dim, time_embed_dim),
-            )
-        else:
-            time_embed_dim = None
-
-        ch = int(conf.channel_mult[0] * conf.model_channels)
-        self.input_blocks = nn.ModuleList([
-            TimestepEmbedSequential(
-                conv_nd(conf.dims, conf.in_channels, ch, 3, padding=1))
-        ])
-        self._feature_size = ch
-        input_block_chans = [ch]
-        ds = 1
-        resolution = conf.image_size
-        for level, mult in enumerate(conf.channel_mult):
-            for _ in range(conf.num_res_blocks):
-                layers = [
-                    ResBlockConfig(
-                        ch,
-                        time_embed_dim,
-                        conf.dropout,
-                        out_channels=int(mult * conf.model_channels),
-                        dims=conf.dims,
-                        use_condition=conf.use_time_condition,
-                        use_checkpoint=conf.use_checkpoint,
-                    ).make_model()
-                ]
-                ch = int(mult * conf.model_channels)
-                if resolution in conf.attention_resolutions:
-                    layers.append(
-                        AttentionBlock(
-                            ch,
-                            use_checkpoint=conf.use_checkpoint,
-                            num_heads=conf.num_heads,
-                            num_head_channels=conf.num_head_channels,
-                            use_new_attention_order=conf.
-                            use_new_attention_order,
-                        ))
-                self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
-                input_block_chans.append(ch)
-            if level != len(conf.channel_mult) - 1:
-                resolution //= 2
-                out_ch = ch
-                self.input_blocks.append(
-                    TimestepEmbedSequential(
-                        ResBlockConfig(
-                            ch,
-                            time_embed_dim,
-                            conf.dropout,
-                            out_channels=out_ch,
-                            dims=conf.dims,
-                            use_condition=conf.use_time_condition,
-                            use_checkpoint=conf.use_checkpoint,
-                            down=True,
-                        ).make_model() if (
-                            conf.resblock_updown
-                        ) else Downsample(ch,
-                                          conf.conv_resample,
-                                          dims=conf.dims,
-                                          out_channels=out_ch)))
-                ch = out_ch
-                input_block_chans.append(ch)
-                ds *= 2
-                self._feature_size += ch
-
-        self.middle_block = TimestepEmbedSequential(
-            ResBlockConfig(
-                ch,
-                time_embed_dim,
-                conf.dropout,
-                dims=conf.dims,
-                use_condition=conf.use_time_condition,
-                use_checkpoint=conf.use_checkpoint,
-            ).make_model(),
-            AttentionBlock(
-                ch,
-                use_checkpoint=conf.use_checkpoint,
-                num_heads=conf.num_heads,
-                num_head_channels=conf.num_head_channels,
-                use_new_attention_order=conf.use_new_attention_order,
-            ),
-            ResBlockConfig(
-                ch,
-                time_embed_dim,
-                conf.dropout,
-                dims=conf.dims,
-                use_condition=conf.use_time_condition,
-                use_checkpoint=conf.use_checkpoint,
-            ).make_model(),
-        )
-        self._feature_size += ch
-        if conf.pool == "adaptivenonzero":
-            self.out = nn.Sequential(
-                normalization(ch),
-                nn.SiLU(),
-                nn.AdaptiveAvgPool2d((1, 1)),
-                conv_nd(conf.dims, ch, conf.out_channels, 1),
-                nn.Flatten(),
-            )
-        else:
-            raise NotImplementedError(f"Unexpected {conf.pool} pooling")
-
-    def forward(self, x, t=None, return_2d_feature=False):
-        """
-        Apply the model to an input batch.
-
-        :param x: an [N x C x ...] Tensor of inputs.
-        :param timesteps: a 1-D batch of timesteps.
-        :return: an [N x K] Tensor of outputs.
-        """
-        if self.conf.use_time_condition:
-            emb = self.time_embed(timestep_embedding(t, self.model_channels))
-        else:
-            emb = None
-
-        results = []
-        h = x.type(self.dtype)
-        for module in self.input_blocks:
-            h = module(h, emb=emb)
-            if self.conf.pool.startswith("spatial"):
-                results.append(h.type(x.dtype).mean(dim=(2, 3)))
-        h = self.middle_block(h, emb=emb)
-        if self.conf.pool.startswith("spatial"):
-            results.append(h.type(x.dtype).mean(dim=(2, 3)))
-            h = th.cat(results, axis=-1)
-        else:
-            h = h.type(x.dtype)
-
-        h_2d = h
-        h = self.out(h)
-
-        if return_2d_feature:
-            return h, h_2d
-        else:
-            return h
-
-    def forward_flatten(self, x):
-        """
-        transform the last 2d feature into a flatten vector
-        """
-        h = self.out(x)
-        return h
-
-
-class SuperResModel(BeatGANsUNetModel):
-    """
-    A UNetModel that performs super-resolution.
-
-    Expects an extra kwarg `low_res` to condition on a low-resolution image.
-    """
-    def __init__(self, image_size, in_channels, *args, **kwargs):
-        super().__init__(image_size, in_channels * 2, *args, **kwargs)
-
-    def forward(self, x, timesteps, low_res=None, **kwargs):
-        _, _, new_height, new_width = x.shape
-        upsampled = F.interpolate(low_res, (new_height, new_width),
-                                  mode="bilinear")
-        x = th.cat([x, upsampled], dim=1)
-        return super().forward(x, timesteps, **kwargs)

code/model/unet_autoenc.py

@@ -1,283 +0,0 @@
-from enum import Enum
-
-import torch
-from torch import Tensor
-from torch.nn.functional import silu
-
-from .latentnet import *
-from .unet import *
-from choices import *
-
-
-@dataclass
-class BeatGANsAutoencConfig(BeatGANsUNetConfig):
-    # number of style channels
-    enc_out_channels: int = 512
-    enc_attn_resolutions: Tuple[int] = None
-    enc_pool: str = 'depthconv'
-    enc_num_res_block: int = 2
-    enc_channel_mult: Tuple[int] = None
-    enc_grad_checkpoint: bool = False
-    latent_net_conf: MLPSkipNetConfig = None
-
-    def make_model(self):
-        return BeatGANsAutoencModel(self)
-
-
-class BeatGANsAutoencModel(BeatGANsUNetModel):
-    def __init__(self, conf: BeatGANsAutoencConfig):
-        super().__init__(conf)
-        self.conf = conf
-
-        # having only time, cond
-        self.time_embed = TimeStyleSeperateEmbed(
-            time_channels=conf.model_channels,
-            time_out_channels=conf.embed_channels,
-        )
-
-        self.encoder = BeatGANsEncoderConfig(
-            image_size=conf.image_size,
-            in_channels=conf.in_channels,
-            model_channels=conf.model_channels,
-            out_hid_channels=conf.enc_out_channels,
-            out_channels=conf.enc_out_channels,
-            num_res_blocks=conf.enc_num_res_block,
-            attention_resolutions=(conf.enc_attn_resolutions
-                                   or conf.attention_resolutions),
-            dropout=conf.dropout,
-            channel_mult=conf.enc_channel_mult or conf.channel_mult,
-            use_time_condition=False,
-            conv_resample=conf.conv_resample,
-            dims=conf.dims,
-            use_checkpoint=conf.use_checkpoint or conf.enc_grad_checkpoint,
-            num_heads=conf.num_heads,
-            num_head_channels=conf.num_head_channels,
-            resblock_updown=conf.resblock_updown,
-            use_new_attention_order=conf.use_new_attention_order,
-            pool=conf.enc_pool,
-        ).make_model()
-
-        if conf.latent_net_conf is not None:
-            self.latent_net = conf.latent_net_conf.make_model()
-
-    def reparameterize(self, mu: Tensor, logvar: Tensor) -> Tensor:
-        """
-        Reparameterization trick to sample from N(mu, var) from
-        N(0,1).
-        :param mu: (Tensor) Mean of the latent Gaussian [B x D]
-        :param logvar: (Tensor) Standard deviation of the latent Gaussian [B x D]
-        :return: (Tensor) [B x D]
-        """
-        assert self.conf.is_stochastic
-        std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std)
-        return eps * std + mu
-
-    def sample_z(self, n: int, device):
-        assert self.conf.is_stochastic
-        return torch.randn(n, self.conf.enc_out_channels, device=device)
-
-    def noise_to_cond(self, noise: Tensor):
-        raise NotImplementedError()
-        assert self.conf.noise_net_conf is not None
-        return self.noise_net.forward(noise)
-
-    def encode(self, x):
-        cond = self.encoder.forward(x)
-        return {'cond': cond}
-
-    @property
-    def stylespace_sizes(self):
-        modules = list(self.input_blocks.modules()) + list(
-            self.middle_block.modules()) + list(self.output_blocks.modules())
-        sizes = []
-        for module in modules:
-            if isinstance(module, ResBlock):
-                linear = module.cond_emb_layers[-1]
-                sizes.append(linear.weight.shape[0])
-        return sizes
-
-    def encode_stylespace(self, x, return_vector: bool = True):
-        """
-        encode to style space
-        """
-        modules = list(self.input_blocks.modules()) + list(
-            self.middle_block.modules()) + list(self.output_blocks.modules())
-        # (n, c)
-        cond = self.encoder.forward(x)
-        S = []
-        for module in modules:
-            if isinstance(module, ResBlock):
-                # (n, c')
-                s = module.cond_emb_layers.forward(cond)
-                S.append(s)
-
-        if return_vector:
-            # (n, sum_c)
-            return torch.cat(S, dim=1)
-        else:
-            return S
-
-    def forward(self,
-                x,
-                t,
-                y=None,
-                x_start=None,
-                cond=None,
-                style=None,
-                noise=None,
-                t_cond=None,
-                **kwargs):
-        """
-        Apply the model to an input batch.
-
-        Args:
-            x_start: the original image to encode
-            cond: output of the encoder
-            noise: random noise (to predict the cond)
-        """
-
-        if t_cond is None:
-            t_cond = t
-
-        if noise is not None:
-            # if the noise is given, we predict the cond from noise
-            cond = self.noise_to_cond(noise)
-
-        if cond is None:
-            if x is not None:
-                assert len(x) == len(x_start), f'{len(x)} != {len(x_start)}'
-
-            tmp = self.encode(x_start)
-            cond = tmp['cond']
-
-        if t is not None:
-            _t_emb = timestep_embedding(t, self.conf.model_channels)
-            _t_cond_emb = timestep_embedding(t_cond, self.conf.model_channels)
-        else:
-            # this happens when training only autoenc
-            _t_emb = None
-            _t_cond_emb = None
-
-        if self.conf.resnet_two_cond:
-            res = self.time_embed.forward(
-                time_emb=_t_emb,
-                cond=cond,
-                time_cond_emb=_t_cond_emb,
-            )
-        else:
-            raise NotImplementedError()
-
-        if self.conf.resnet_two_cond:
-            # two cond: first = time emb, second = cond_emb
-            emb = res.time_emb
-            cond_emb = res.emb
-        else:
-            # one cond = combined of both time and cond
-            emb = res.emb
-            cond_emb = None
-
-        # override the style if given
-        style = style or res.style
-
-        assert (y is not None) == (
-            self.conf.num_classes is not None
-        ), "must specify y if and only if the model is class-conditional"
-
-        if self.conf.num_classes is not None:
-            raise NotImplementedError()
-            # assert y.shape == (x.shape[0], )
-            # emb = emb + self.label_emb(y)
-
-        # where in the model to supply time conditions
-        enc_time_emb = emb
-        mid_time_emb = emb
-        dec_time_emb = emb
-        # where in the model to supply style conditions
-        enc_cond_emb = cond_emb
-        mid_cond_emb = cond_emb
-        dec_cond_emb = cond_emb
-
-        # hs = []
-        hs = [[] for _ in range(len(self.conf.channel_mult))]
-
-        if x is not None:
-            h = x.type(self.dtype)
-
-            # input blocks
-            k = 0
-            for i in range(len(self.input_num_blocks)):
-                for j in range(self.input_num_blocks[i]):
-                    h = self.input_blocks[k](h,
-                                             emb=enc_time_emb,
-                                             cond=enc_cond_emb)
-
-                    # print(i, j, h.shape)
-                    hs[i].append(h)
-                    k += 1
-            assert k == len(self.input_blocks)
-
-            # middle blocks
-            h = self.middle_block(h, emb=mid_time_emb, cond=mid_cond_emb)
-        else:
-            # no lateral connections
-            # happens when training only the autonecoder
-            h = None
-            hs = [[] for _ in range(len(self.conf.channel_mult))]
-
-        # output blocks
-        k = 0
-        for i in range(len(self.output_num_blocks)):
-            for j in range(self.output_num_blocks[i]):
-                # take the lateral connection from the same layer (in reserve)
-                # until there is no more, use None
-                try:
-                    lateral = hs[-i - 1].pop()
-                    # print(i, j, lateral.shape)
-                except IndexError:
-                    lateral = None
-                    # print(i, j, lateral)
-
-                h = self.output_blocks[k](h,
-                                          emb=dec_time_emb,
-                                          cond=dec_cond_emb,
-                                          lateral=lateral)
-                k += 1
-
-        pred = self.out(h)
-        return AutoencReturn(pred=pred, cond=cond)
-
-
-class AutoencReturn(NamedTuple):
-    pred: Tensor
-    cond: Tensor = None
-
-
-class EmbedReturn(NamedTuple):
-    # style and time
-    emb: Tensor = None
-    # time only
-    time_emb: Tensor = None
-    # style only (but could depend on time)
-    style: Tensor = None
-
-
-class TimeStyleSeperateEmbed(nn.Module):
-    # embed only style
-    def __init__(self, time_channels, time_out_channels):
-        super().__init__()
-        self.time_embed = nn.Sequential(
-            linear(time_channels, time_out_channels),
-            nn.SiLU(),
-            linear(time_out_channels, time_out_channels),
-        )
-        self.style = nn.Identity()
-
-    def forward(self, time_emb=None, cond=None, **kwargs):
-        if time_emb is None:
-            # happens with autoenc training mode
-            time_emb = None
-        else:
-            time_emb = self.time_embed(time_emb)
-        style = self.style(cond)
-        return EmbedReturn(emb=style, time_emb=time_emb, style=style)

code/networks/discriminator.py

@@ -1,259 +0,0 @@
-import math
-import torch
-from torch.nn import functional as F
-from torch import nn
-
-
-def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return F.leaky_relu(input + bias, negative_slope) * scale
-
-
-class FusedLeakyReLU(nn.Module):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-        self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        # print("FusedLeakyReLU: ", input.abs().mean())
-        out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-        # print("FusedLeakyReLU: ", out.abs().mean())
-        return out
-
-
-def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
-    _, minor, in_h, in_w = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, minor, in_h, 1, in_w, 1)
-    out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
-    out = out.view(-1, minor, in_h * up_y, in_w * up_x)
-
-    out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
-    out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0), max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
-
-    # out = out.permute(0, 3, 1, 2)
-    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-                      in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
-    # out = out.permute(0, 2, 3, 1)
-
-    return out[:, :, ::down_y, ::down_x]
-
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
-
-
-def make_kernel(k):
-    k = torch.tensor(k, dtype=torch.float32)
-
-    if k.ndim == 1:
-        k = k[None, :] * k[:, None]
-
-    k /= k.sum()
-
-    return k
-
-
-class Blur(nn.Module):
-    def __init__(self, kernel, pad, upsample_factor=1):
-        super().__init__()
-
-        kernel = make_kernel(kernel)
-
-        if upsample_factor > 1:
-            kernel = kernel * (upsample_factor ** 2)
-
-        self.register_buffer('kernel', kernel)
-
-        self.pad = pad
-
-    def forward(self, input):
-        return upfirdn2d(input, self.kernel, pad=self.pad)
-
-
-class ScaledLeakyReLU(nn.Module):
-    def __init__(self, negative_slope=0.2):
-        super().__init__()
-
-        self.negative_slope = negative_slope
-
-    def forward(self, input):
-        return F.leaky_relu(input, negative_slope=self.negative_slope)
-
-
-class EqualConv2d(nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
-        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
-
-        self.stride = stride
-        self.padding = padding
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_channel))
-        else:
-            self.bias = None
-
-    def forward(self, input):
-
-        return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride,
-                        padding=self.padding, )
-
-    def __repr__(self):
-        return (
-            f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
-            f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
-        )
-
-
-class EqualLinear(nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
-        else:
-            self.bias = None
-
-        self.activation = activation
-
-        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
-        self.lr_mul = lr_mul
-
-    def forward(self, input):
-
-        if self.activation:
-            out = F.linear(input, self.weight * self.scale)
-            out = fused_leaky_relu(out, self.bias * self.lr_mul)
-        else:
-            out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
-
-        return out
-
-    def __repr__(self):
-        return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
-
-
-class ConvLayer(nn.Sequential):
-    def __init__(
-            self,
-            in_channel,
-            out_channel,
-            kernel_size,
-            downsample=False,
-            blur_kernel=[1, 3, 3, 1],
-            bias=True,
-            activate=True,
-    ):
-        layers = []
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
-
-            stride = 2
-            self.padding = 0
-
-        else:
-            stride = 1
-            self.padding = kernel_size // 2
-
-        layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
-                                  bias=bias and not activate))
-
-        if activate:
-            if bias:
-                layers.append(FusedLeakyReLU(out_channel))
-            else:
-                layers.append(ScaledLeakyReLU(0.2))
-
-        super().__init__(*layers)
-
-
-class ResBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.conv1 = ConvLayer(in_channel, in_channel, 3)
-        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
-
-        self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
-
-    def forward(self, input):
-        out = self.conv1(input)
-        out = self.conv2(out)
-
-        skip = self.skip(input)
-        out = (out + skip) / math.sqrt(2)
-
-        return out
-
-
-class Discriminator(nn.Module):
-    def __init__(self, size, channel_multiplier=1, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.size = size
-
-        channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        convs = [ConvLayer(3, channels[size], 1)]
-        log_size = int(math.log(size, 2))
-        in_channel = channels[size]
-
-        for i in range(log_size, 2, -1):
-            out_channel = channels[2 ** (i - 1)]
-            convs.append(ResBlock(in_channel, out_channel, blur_kernel))
-            in_channel = out_channel
-
-        self.convs = nn.Sequential(*convs)
-
-        self.stddev_group = 4
-        self.stddev_feat = 1
-
-        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
-        self.final_linear = nn.Sequential(
-            EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'),
-            EqualLinear(channels[4], 1),
-        )
-
-    def forward(self, input):
-        out = self.convs(input)
-        batch, channel, height, width = out.shape
-
-        group = min(batch, self.stddev_group)
-        stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
-        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
-        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
-        stddev = stddev.repeat(group, 1, height, width)
-        out = torch.cat([out, stddev], 1)
-
-        out = self.final_conv(out)
-
-        out = out.view(batch, -1)
-        out = self.final_linear(out)
-
-        return out

code/networks/encoder.py

@@ -1,374 +0,0 @@
-import math
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return F.leaky_relu(input + bias, negative_slope) * scale
-
-class FusedLeakyReLU(nn.Module):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-        self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-        return out
-
-
-def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
-    _, minor, in_h, in_w = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, minor, in_h, 1, in_w, 1)
-    out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
-    out = out.view(-1, minor, in_h * up_y, in_w * up_x)
-
-    out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
-    out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0),
-          max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
-
-    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-                      in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
-
-    return out[:, :, ::down_y, ::down_x]
-
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
-
-
-def make_kernel(k):
-    k = torch.tensor(k, dtype=torch.float32)
-
-    if k.ndim == 1:
-        k = k[None, :] * k[:, None]
-
-    k /= k.sum()
-
-    return k
-
-
-class Blur(nn.Module):
-    def __init__(self, kernel, pad, upsample_factor=1):
-        super().__init__()
-
-        kernel = make_kernel(kernel)
-
-        if upsample_factor > 1:
-            kernel = kernel * (upsample_factor ** 2)
-
-        self.register_buffer('kernel', kernel)
-
-        self.pad = pad
-
-    def forward(self, input):
-        return upfirdn2d(input, self.kernel, pad=self.pad)
-
-
-class ScaledLeakyReLU(nn.Module):
-    def __init__(self, negative_slope=0.2):
-        super().__init__()
-
-        self.negative_slope = negative_slope
-
-    def forward(self, input):
-        return F.leaky_relu(input, negative_slope=self.negative_slope)
-
-
-class EqualConv2d(nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
-        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
-
-        self.stride = stride
-        self.padding = padding
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_channel))
-        else:
-            self.bias = None
-
-    def forward(self, input):
-
-        return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding)
-
-    def __repr__(self):
-        return (
-            f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
-            f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
-        )
-
-
-class EqualLinear(nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
-        else:
-            self.bias = None
-
-        self.activation = activation
-
-        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
-        self.lr_mul = lr_mul
-
-    def forward(self, input):
-
-        if self.activation:
-            out = F.linear(input, self.weight * self.scale)
-            out = fused_leaky_relu(out, self.bias * self.lr_mul)
-        else:
-            out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
-
-        return out
-
-    def __repr__(self):
-        return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
-
-
-class ConvLayer(nn.Sequential):
-    def __init__(
-            self,
-            in_channel,
-            out_channel,
-            kernel_size,
-            downsample=False,
-            blur_kernel=[1, 3, 3, 1],
-            bias=True,
-            activate=True,
-    ):
-        layers = []
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
-
-            stride = 2
-            self.padding = 0
-
-        else:
-            stride = 1
-            self.padding = kernel_size // 2
-
-        layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
-                                  bias=bias and not activate))
-
-        if activate:
-            if bias:
-                layers.append(FusedLeakyReLU(out_channel))
-            else:
-                layers.append(ScaledLeakyReLU(0.2))
-
-        super().__init__(*layers)
-
-
-class ResBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.conv1 = ConvLayer(in_channel, in_channel, 3)
-        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
-
-        self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
-
-    def forward(self, input):
-        out = self.conv1(input)
-        out = self.conv2(out)
-
-        skip = self.skip(input)
-        out = (out + skip) / math.sqrt(2)
-
-        return out
-
-class WeightedSumLayer(nn.Module):
-    def __init__(self, num_tensors=8):
-        super(WeightedSumLayer, self).__init__()
-
-        self.weights = nn.Parameter(torch.randn(num_tensors))
-    
-    def forward(self, tensor_list):
-
-        weights = torch.softmax(self.weights, dim=0)
-        weighted_sum = torch.zeros_like(tensor_list[0])
-        for tensor, weight in zip(tensor_list, weights):
-            weighted_sum += tensor * weight
-        
-        return weighted_sum
-
-class EncoderApp(nn.Module):
-    def __init__(self, size, w_dim=512, fusion_type=''):
-        super(EncoderApp, self).__init__()
-
-        channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256,
-            128: 128,
-            256: 64,
-            512: 32,
-            1024: 16
-        }
-
-        self.w_dim = w_dim
-        log_size = int(math.log(size, 2))
-
-        self.convs = nn.ModuleList()
-        self.convs.append(ConvLayer(3, channels[size], 1))
-
-        in_channel = channels[size]
-        for i in range(log_size, 2, -1):
-            out_channel = channels[2 ** (i - 1)]
-            self.convs.append(ResBlock(in_channel, out_channel))
-            in_channel = out_channel
-
-        self.convs.append(EqualConv2d(in_channel, self.w_dim, 4, padding=0, bias=False))
-        
-        self.fusion_type = fusion_type
-        assert self.fusion_type == 'weighted_sum'
-        if self.fusion_type == 'weighted_sum':
-            print(f'HAL layer is enabled!')
-            self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
-            self.fc1 = EqualLinear(64, 512)
-            self.fc2 = EqualLinear(128, 512)
-            self.fc3 = EqualLinear(256, 512)
-            self.ws = WeightedSumLayer()
-
-    def forward(self, x):
-
-        res = []
-        h = x
-        pooled_h_lists = []
-        for i, conv in enumerate(self.convs):
-            h = conv(h)
-            if self.fusion_type == 'weighted_sum':
-                pooled_h = self.adaptive_pool(h).view(x.size(0), -1) 
-                if i == 0:
-                    pooled_h_lists.append(self.fc1(pooled_h))
-                elif i == 1:
-                    pooled_h_lists.append(self.fc2(pooled_h))
-                elif i == 2:
-                    pooled_h_lists.append(self.fc3(pooled_h))
-                else:
-                    pooled_h_lists.append(pooled_h)
-            res.append(h)
-        
-        if self.fusion_type == 'weighted_sum':
-            last_layer = self.ws(pooled_h_lists)
-        else:  
-            last_layer = res[-1].squeeze(-1).squeeze(-1)
-        layer_features = res[::-1][2:]
-        
-        return last_layer, layer_features
-
-
-class DecouplingModel(nn.Module):
-    def __init__(self, input_dim, hidden_dim, output_dim):
-        super(DecouplingModel, self).__init__()
-        
-        # identity_excluded_net is called identity encoder in the paper
-        self.identity_net = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, output_dim)
-        )
-        
-        self.identity_net_density = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, output_dim)
-        )
-        
-        # identity_excluded_net is called motion encoder in the paper
-        self.identity_excluded_net = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, output_dim)
-        )
-
-    def forward(self, x):
-
-        id_, id_rm =  self.identity_net(x), self.identity_excluded_net(x)
-        id_density = self.identity_net_density(id_)
-        return id_, id_rm, id_density
-
-class Encoder(nn.Module):
-    def __init__(self, size, dim=512, dim_motion=20, weighted_sum=False):
-        super(Encoder, self).__init__()
-
-        # image encoder
-        self.net_app = EncoderApp(size, dim, weighted_sum)
-        
-        # decouping network
-        self.net_decouping = DecouplingModel(dim, dim, dim)
-
-        # part of the motion encoder
-        fc = [EqualLinear(dim, dim)]
-        for i in range(3):
-            fc.append(EqualLinear(dim, dim))
-
-        fc.append(EqualLinear(dim, dim_motion))
-        self.fc = nn.Sequential(*fc)
-
-    def enc_app(self, x):
-
-        h_source = self.net_app(x)
-
-        return h_source
-
-    def enc_motion(self, x):
-
-        h, _ = self.net_app(x)
-        h_motion = self.fc(h)
-
-        return h_motion
-    
-    def encode_image_obj(self, image_obj):
-        feat, _ = self.net_app(image_obj)
-        id_emb, idrm_emb, id_density_emb = self.net_decouping(feat)
-        return id_emb, idrm_emb, id_density_emb
-
-    def forward(self, input_source, input_target, input_face, input_aug):
-
-
-        if input_target is not None:
-
-            h_source, feats = self.net_app(input_source)
-            h_target, _ = self.net_app(input_target)
-            h_face, _ = self.net_app(input_face)
-            h_aug, _ = self.net_app(input_aug)
-            
-            h_source_id_emb, h_source_idrm_emb, h_source_id_density_emb = self.net_decouping(h_source)
-            h_target_id_emb, h_target_idrm_emb, h_target_id_density_emb = self.net_decouping(h_target)
-            h_face_id_emb, h_face_idrm_emb, h_face_id_density_emb = self.net_decouping(h_face)
-            h_aug_id_emb, h_aug_idrm_emb, h_aug_id_density_emb = self.net_decouping(h_aug)
-
-            h_target_motion_target = self.fc(h_target_idrm_emb)
-            h_another_face_target =  self.fc(h_face_idrm_emb)
-            
-        else:
-            h_source, feats = self.net_app(input_source)
-
-
-        return {'h_source':h_source, 'h_motion':h_target_motion_target, 'feats':feats, 'h_another_face_target':h_another_face_target, 'h_face':h_face, \
-                'h_source_id_emb':h_source_id_emb, 'h_source_idrm_emb':h_source_idrm_emb,  'h_source_id_density_emb':h_source_id_density_emb, \
-                'h_target_id_emb':h_target_id_emb, 'h_target_idrm_emb':h_target_idrm_emb,  'h_target_id_density_emb':h_target_id_density_emb, \
-                'h_face_id_emb':h_face_id_emb, 'h_face_idrm_emb':h_face_idrm_emb, 'h_face_id_density_emb':h_face_id_density_emb, \
-                'h_aug_id_emb':h_aug_id_emb, 'h_aug_idrm_emb':h_aug_idrm_emb ,'h_aug_id_density_emb':h_aug_id_density_emb, \
-                }

code/networks/generator.py

@@ -1,27 +0,0 @@
-from torch import nn
-from .encoder import Encoder
-from .styledecoder import Synthesis
-
-
-class Generator(nn.Module):
-    def __init__(self, size, style_dim=512, motion_dim=20, channel_multiplier=1, blur_kernel=[1, 3, 3, 1]):
-        super(Generator, self).__init__()
-
-        # encoder
-        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 = 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)
-        # import pdb;pdb.set_trace()
-        img_recon = self.dec(wa, alpha, feats)
-
-        return img_recon

code/networks/styledecoder.py

@@ -1,527 +0,0 @@
-import math
-import torch
-from torch import nn
-from torch.nn import functional as F
-import numpy as np
-
-
-def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return F.leaky_relu(input + bias, negative_slope) * scale
-
-
-class FusedLeakyReLU(nn.Module):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-        self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-        return out
-
-
-def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
-    _, minor, in_h, in_w = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, minor, in_h, 1, in_w, 1)
-    out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
-    out = out.view(-1, minor, in_h * up_y, in_w * up_x)
-
-    out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
-    out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0),
-          max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
-
-    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-                      in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
-    return out[:, :, ::down_y, ::down_x]
-
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
-
-
-class PixelNorm(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, input):
-        return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
-
-
-class MotionPixelNorm(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, input):
-        return input * torch.rsqrt(torch.mean(input ** 2, dim=2, keepdim=True) + 1e-8)
-
-
-def make_kernel(k):
-    k = torch.tensor(k, dtype=torch.float32)
-
-    if k.ndim == 1:
-        k = k[None, :] * k[:, None]
-
-    k /= k.sum()
-
-    return k
-
-
-class Upsample(nn.Module):
-    def __init__(self, kernel, factor=2):
-        super().__init__()
-
-        self.factor = factor
-        kernel = make_kernel(kernel) * (factor ** 2)
-        self.register_buffer('kernel', kernel)
-
-        p = kernel.shape[0] - factor
-
-        pad0 = (p + 1) // 2 + factor - 1
-        pad1 = p // 2
-
-        self.pad = (pad0, pad1)
-
-    def forward(self, input):
-        return upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
-
-
-class Downsample(nn.Module):
-    def __init__(self, kernel, factor=2):
-        super().__init__()
-
-        self.factor = factor
-        kernel = make_kernel(kernel)
-        self.register_buffer('kernel', kernel)
-
-        p = kernel.shape[0] - factor
-
-        pad0 = (p + 1) // 2
-        pad1 = p // 2
-
-        self.pad = (pad0, pad1)
-
-    def forward(self, input):
-        return upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
-
-
-class Blur(nn.Module):
-    def __init__(self, kernel, pad, upsample_factor=1):
-        super().__init__()
-
-        kernel = make_kernel(kernel)
-
-        if upsample_factor > 1:
-            kernel = kernel * (upsample_factor ** 2)
-
-        self.register_buffer('kernel', kernel)
-
-        self.pad = pad
-
-    def forward(self, input):
-        return upfirdn2d(input, self.kernel, pad=self.pad)
-
-
-class EqualConv2d(nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
-        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
-
-        self.stride = stride
-        self.padding = padding
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_channel))
-        else:
-            self.bias = None
-
-    def forward(self, input):
-
-        return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding, )
-
-    def __repr__(self):
-        return (
-            f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
-            f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
-        )
-
-
-class EqualLinear(nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
-        else:
-            self.bias = None
-
-        self.activation = activation
-
-        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
-        self.lr_mul = lr_mul
-
-    def forward(self, input):
-
-        if self.activation:
-            out = F.linear(input, self.weight * self.scale)
-            out = fused_leaky_relu(out, self.bias * self.lr_mul)
-        else:
-            out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
-
-        return out
-
-    def __repr__(self):
-        return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
-
-
-class ScaledLeakyReLU(nn.Module):
-    def __init__(self, negative_slope=0.2):
-        super().__init__()
-
-        self.negative_slope = negative_slope
-
-    def forward(self, input):
-        return F.leaky_relu(input, negative_slope=self.negative_slope)
-
-
-class ModulatedConv2d(nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, style_dim, demodulate=True, upsample=False,
-                 downsample=False, blur_kernel=[1, 3, 3, 1], ):
-        super().__init__()
-
-        self.eps = 1e-8
-        self.kernel_size = kernel_size
-        self.in_channel = in_channel
-        self.out_channel = out_channel
-        self.upsample = upsample
-        self.downsample = downsample
-
-        if upsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) - (kernel_size - 1)
-            pad0 = (p + 1) // 2 + factor - 1
-            pad1 = p // 2 + 1
-
-            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            self.blur = Blur(blur_kernel, pad=(pad0, pad1))
-
-        fan_in = in_channel * kernel_size ** 2
-        self.scale = 1 / math.sqrt(fan_in)
-        self.padding = kernel_size // 2
-
-        self.weight = nn.Parameter(torch.randn(1, out_channel, in_channel, kernel_size, kernel_size))
-
-        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
-        self.demodulate = demodulate
-
-    def __repr__(self):
-        return (
-            f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, '
-            f'upsample={self.upsample}, downsample={self.downsample})'
-        )
-
-    def forward(self, input, style):
-        batch, in_channel, height, width = input.shape
-
-        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
-        weight = self.scale * self.weight * style
-
-        if self.demodulate:
-            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
-            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
-
-        weight = weight.view(batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size)
-
-        if self.upsample:
-            input = input.view(1, batch * in_channel, height, width)
-            weight = weight.view(batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size)
-            weight = weight.transpose(1, 2).reshape(batch * in_channel, self.out_channel, self.kernel_size,
-                                                    self.kernel_size)
-            out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch)
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-            out = self.blur(out)
-        elif self.downsample:
-            input = self.blur(input)
-            _, _, height, width = input.shape
-            input = input.view(1, batch * in_channel, height, width)
-            out = F.conv2d(input, weight, padding=0, stride=2, groups=batch)
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-        else:
-            input = input.view(1, batch * in_channel, height, width)
-            out = F.conv2d(input, weight, padding=self.padding, groups=batch)
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-
-        return out
-
-
-class NoiseInjection(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.zeros(1))
-
-    def forward(self, image, noise=None):
-
-        if noise is None:
-            return image
-        else:
-            return image + self.weight * noise
-
-
-class ConstantInput(nn.Module):
-    def __init__(self, channel, size=4):
-        super().__init__()
-
-        self.input = nn.Parameter(torch.randn(1, channel, size, size))
-
-    def forward(self, input):
-        batch = input.shape[0]
-        out = self.input.repeat(batch, 1, 1, 1)
-
-        return out
-
-
-class StyledConv(nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, style_dim, upsample=False, blur_kernel=[1, 3, 3, 1],
-                 demodulate=True):
-        super().__init__()
-
-        self.conv = ModulatedConv2d(
-            in_channel,
-            out_channel,
-            kernel_size,
-            style_dim,
-            upsample=upsample,
-            blur_kernel=blur_kernel,
-            demodulate=demodulate,
-        )
-
-        self.noise = NoiseInjection()
-        self.activate = FusedLeakyReLU(out_channel)
-
-    def forward(self, input, style, noise=None):
-        out = self.conv(input, style)
-        out = self.noise(out, noise=noise)
-        out = self.activate(out)
-
-        return out
-
-
-class ConvLayer(nn.Sequential):
-    def __init__(
-            self,
-            in_channel,
-            out_channel,
-            kernel_size,
-            downsample=False,
-            blur_kernel=[1, 3, 3, 1],
-            bias=True,
-            activate=True,
-    ):
-        layers = []
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
-
-            stride = 2
-            self.padding = 0
-
-        else:
-            stride = 1
-            self.padding = kernel_size // 2
-
-        layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
-                                  bias=bias and not activate))
-
-        if activate:
-            if bias:
-                layers.append(FusedLeakyReLU(out_channel))
-            else:
-                layers.append(ScaledLeakyReLU(0.2))
-
-        super().__init__(*layers)
-
-
-class ToRGB(nn.Module):
-    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        if upsample:
-            self.upsample = Upsample(blur_kernel)
-
-        self.conv = ConvLayer(in_channel, 3, 1)
-        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
-
-    def forward(self, input, skip=None):
-        out = self.conv(input)
-        out = out + self.bias
-
-        if skip is not None:
-            skip = self.upsample(skip)
-            out = out + skip
-
-        return out
-
-
-class ToFlow(nn.Module):
-    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        if upsample:
-            self.upsample = Upsample(blur_kernel)
-        
-        self.style_dim = style_dim
-        self.in_channel = in_channel
-        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
-        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
-
-    def forward(self, input, style, feat, skip=None): # input 是来自上一层的 feature， style 是 512 的 condition， feat 是来自于 unet 的跳层
-        out = self.conv(input, style)
-        out = out + self.bias
-
-        # warping
-        xs = np.linspace(-1, 1, input.size(2))
-        
-        xs = np.meshgrid(xs, xs)
-        xs = np.stack(xs, 2)
-
-        xs = torch.tensor(xs, requires_grad=False).float().unsqueeze(0).repeat(input.size(0), 1, 1, 1).to(input.device)
-        # import pdb;pdb.set_trace()
-        if skip is not None:
-            skip = self.upsample(skip)
-            out = out + skip
-
-        sampler = torch.tanh(out[:, 0:2, :, :])
-        mask = torch.sigmoid(out[:, 2:3, :, :])
-        flow = sampler.permute(0, 2, 3, 1) + xs # xs在这里相当于一个 location 的位置
-
-        feat_warp = F.grid_sample(feat, flow) * mask
-        # import pdb;pdb.set_trace()
-        return feat_warp, feat_warp + input * (1.0 - mask), out
-
-
-class Direction(nn.Module):
-    def __init__(self, motion_dim):
-        super(Direction, self).__init__()
-
-        self.weight = nn.Parameter(torch.randn(512, motion_dim))
-
-    def forward(self, input):
-        # input: (bs*t) x 512
-
-        weight = self.weight + 1e-8
-        Q, R = torch.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]
-
-        if input is None:
-            return Q
-        else:
-            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
-            out = torch.matmul(input_diag, Q.T)
-            out = torch.sum(out, dim=1)
-
-            return out
-
-class Synthesis(nn.Module):
-    def __init__(self, size, style_dim, motion_dim, blur_kernel=[1, 3, 3, 1], channel_multiplier=1):
-        super(Synthesis, self).__init__()
-
-        self.size = size
-        self.style_dim = style_dim
-        self.motion_dim = motion_dim
-
-        self.direction = Direction(motion_dim) # Linear Motion Decomposition (LMD) from LIA
-
-        self.channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        self.input = ConstantInput(self.channels[4])
-        self.conv1 = StyledConv(self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel)
-        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
-
-        self.log_size = int(math.log(size, 2))
-        self.num_layers = (self.log_size - 2) * 2 + 1
-
-        self.convs = nn.ModuleList()
-        self.upsamples = nn.ModuleList()
-        self.to_rgbs = nn.ModuleList()
-        self.to_flows = nn.ModuleList()
-
-        in_channel = self.channels[4]
-
-        for i in range(3, self.log_size + 1):
-            out_channel = self.channels[2 ** i]
-
-            self.convs.append(StyledConv(in_channel, out_channel, 3, style_dim, upsample=True,
-                                         blur_kernel=blur_kernel))
-            self.convs.append(StyledConv(out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel))
-            self.to_rgbs.append(ToRGB(out_channel, style_dim))
-
-            self.to_flows.append(ToFlow(out_channel, style_dim))
-
-            in_channel = out_channel
-
-        self.n_latent = self.log_size * 2 - 2
-        
-    def forward(self, source_before_decoupling, target_motion, feats):
-
-        directions = self.direction(target_motion)
-        latent = source_before_decoupling + directions  # wa + directions
-
-        inject_index = self.n_latent
-        latent = latent.unsqueeze(1).repeat(1, inject_index, 1)
-        
-        out = self.input(latent)
-        out = self.conv1(out, latent[:, 0])
-
-        i = 1
-        for conv1, conv2, to_rgb, to_flow, feat in zip(self.convs[::2], self.convs[1::2], self.to_rgbs,
-                                                       self.to_flows, feats):
-            out = conv1(out, latent[:, i])
-            out = conv2(out, latent[:, i + 1])
-            if out.size(2) == 8:
-                out_warp, out, skip_flow = to_flow(out, latent[:, i + 2], feat)
-                skip = to_rgb(out_warp)
-            else:
-                out_warp, out, skip_flow = to_flow(out, latent[:, i + 2], feat, skip_flow)
-                skip = to_rgb(out_warp, skip)
-            i += 2
-
-        img = skip
-
-        return img
-        
-        
-

code/networks/utils.py

@@ -1,53 +0,0 @@
-from torch import nn
-import torch.nn.functional as F
-import torch
-
-
-class AntiAliasInterpolation2d(nn.Module):
-    """
-    Band-limited downsampling, for better preservation of the input signal.
-    """
-
-    def __init__(self, channels, scale):
-        super(AntiAliasInterpolation2d, self).__init__()
-        sigma = (1 / scale - 1) / 2
-        kernel_size = 2 * round(sigma * 4) + 1
-        self.ka = kernel_size // 2
-        self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka
-
-        kernel_size = [kernel_size, kernel_size]
-        sigma = [sigma, sigma]
-        # The gaussian kernel is the product of the
-        # gaussian function of each dimension.
-        kernel = 1
-        meshgrids = torch.meshgrid(
-            [
-                torch.arange(size, dtype=torch.float32)
-                for size in kernel_size
-            ]
-        )
-        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
-            mean = (size - 1) / 2
-            kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2))
-
-        # Make sure sum of values in gaussian kernel equals 1.
-        kernel = kernel / torch.sum(kernel)
-        # Reshape to depthwise convolutional weight
-        kernel = kernel.view(1, 1, *kernel.size())
-        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
-
-        self.register_buffer('weight', kernel)
-        self.groups = channels
-        self.scale = scale
-        inv_scale = 1 / scale
-        self.int_inv_scale = int(inv_scale)
-
-    def forward(self, input):
-        if self.scale == 1.0:
-            return input
-
-        out = F.pad(input, (self.ka, self.kb, self.ka, self.kb))
-        out = F.conv2d(out, weight=self.weight, groups=self.groups)
-        out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale]
-
-        return out

code/renderer.py

@@ -1,25 +0,0 @@
-from config import *
-
-def render_condition(
-    conf: TrainConfig,
-    model,
-    sampler, start, motion_direction_start, audio_driven, \
-        face_location, face_scale, \
-        yaw_pitch_roll, noisyT, control_flag,
-):
-    if conf.train_mode == TrainMode.diffusion:
-        assert conf.model_type.has_autoenc()
-
-        return sampler.sample(model=model,
-                              noise=noisyT,
-                              model_kwargs={
-                                  'motion_direction_start': motion_direction_start, 
-                                  'yaw_pitch_roll': yaw_pitch_roll,
-                                  'start': start,
-                                  'audio_driven': audio_driven,
-                                  'face_location': face_location,
-                                  'face_scale': face_scale,
-                                  'control_flag': control_flag
-                              })
-    else:
-        raise NotImplementedError()

code/templates.py

@@ -1,301 +0,0 @@
-from experiment import *
-
-
-def ddpm():
-    """
-    base configuration for all DDIM-based models.
-    """
-    conf = TrainConfig()
-    conf.batch_size = 32
-    conf.beatgans_gen_type = GenerativeType.ddim
-    conf.beta_scheduler = 'linear'
-    conf.data_name = 'ffhq'
-    conf.diffusion_type = 'beatgans'
-    conf.eval_ema_every_samples = 200_000
-    conf.eval_every_samples = 200_000
-    conf.fp16 = True
-    conf.lr = 1e-4
-    conf.model_name = ModelName.beatgans_ddpm
-    conf.net_attn = (16, )
-    conf.net_beatgans_attn_head = 1
-    conf.net_beatgans_embed_channels = 512
-    conf.net_ch_mult = (1, 2, 4, 8)
-    conf.net_ch = 64
-    conf.sample_size = 32
-    conf.T_eval = 20
-    conf.T = 1000
-    conf.make_model_conf()
-    return conf
-
-
-def autoenc_base():
-    """
-    base configuration for all Diff-AE models.
-    """
-    conf = TrainConfig()
-    conf.batch_size = 32
-    conf.beatgans_gen_type = GenerativeType.ddim
-    conf.beta_scheduler = 'linear'
-    conf.data_name = 'ffhq'
-    conf.diffusion_type = 'beatgans'
-    conf.eval_ema_every_samples = 200_000
-    conf.eval_every_samples = 200_000
-    conf.fp16 = True
-    conf.lr = 1e-4
-    conf.model_name = ModelName.beatgans_autoenc
-    conf.net_attn = (16, )
-    conf.net_beatgans_attn_head = 1
-    conf.net_beatgans_embed_channels = 512
-    conf.net_beatgans_resnet_two_cond = True
-    conf.net_ch_mult = (1, 2, 4, 8)
-    conf.net_ch = 64
-    conf.net_enc_channel_mult = (1, 2, 4, 8, 8)
-    conf.net_enc_pool = 'adaptivenonzero'
-    conf.sample_size = 32
-    conf.T_eval = 20
-    conf.T = 1000
-    conf.make_model_conf()
-    return conf
-
-
-def ffhq64_ddpm():
-    conf = ddpm()
-    conf.data_name = 'ffhqlmdb256'
-    conf.warmup = 0
-    conf.total_samples = 72_000_000
-    conf.scale_up_gpus(4)
-    return conf
-
-
-def ffhq64_autoenc():
-    conf = autoenc_base()
-    conf.data_name = 'ffhqlmdb256'
-    conf.warmup = 0
-    conf.total_samples = 72_000_000
-    conf.net_ch_mult = (1, 2, 4, 8)
-    conf.net_enc_channel_mult = (1, 2, 4, 8, 8)
-    conf.eval_every_samples = 1_000_000
-    conf.eval_ema_every_samples = 1_000_000
-    conf.scale_up_gpus(4)
-    conf.make_model_conf()
-    return conf
-
-
-def celeba64d2c_ddpm():
-    conf = ffhq128_ddpm()
-    conf.data_name = 'celebalmdb'
-    conf.eval_every_samples = 10_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.total_samples = 72_000_000
-    conf.name = 'celeba64d2c_ddpm'
-    return conf
-
-
-def celeba64d2c_autoenc():
-    conf = ffhq64_autoenc()
-    conf.data_name = 'celebalmdb'
-    conf.eval_every_samples = 10_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.total_samples = 72_000_000
-    conf.name = 'celeba64d2c_autoenc'
-    return conf
-
-
-def ffhq128_ddpm():
-    conf = ddpm()
-    conf.data_name = 'ffhqlmdb256'
-    conf.warmup = 0
-    conf.total_samples = 48_000_000
-    conf.img_size = 128
-    conf.net_ch = 128
-    # channels:
-    # 3 => 128 * 1 => 128 * 1 => 128 * 2 => 128 * 3 => 128 * 4
-    # sizes:
-    # 128 => 128 => 64 => 32 => 16 => 8
-    conf.net_ch_mult = (1, 1, 2, 3, 4)
-    conf.eval_every_samples = 1_000_000
-    conf.eval_ema_every_samples = 1_000_000
-    conf.scale_up_gpus(4)
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.make_model_conf()
-    return conf
-
-
-def ffhq128_autoenc_base():
-    conf = autoenc_base()
-    conf.data_name = 'ffhqlmdb256'
-    conf.scale_up_gpus(4)
-    conf.img_size = 128
-    conf.net_ch = 128
-    # final resolution = 8x8
-    conf.net_ch_mult = (1, 1, 2, 3, 4)
-    # final resolution = 4x4
-    conf.net_enc_channel_mult = (1, 1, 2, 3, 4, 4)
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.make_model_conf()
-    return conf
-
-
-def ffhq256_autoenc():
-    conf = ffhq128_autoenc_base()
-    conf.img_size = 256
-    conf.net_ch = 128
-    conf.net_ch_mult = (1, 1, 2, 2, 4, 4)
-    conf.net_enc_channel_mult = (1, 1, 2, 2, 4, 4, 4)
-    conf.eval_every_samples = 10_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.total_samples = 200_000_000
-    conf.batch_size = 64
-    conf.make_model_conf()
-    conf.name = 'ffhq256_autoenc'
-    return conf
-
-
-def ffhq256_autoenc_eco():
-    conf = ffhq128_autoenc_base()
-    conf.img_size = 256
-    conf.net_ch = 128
-    conf.net_ch_mult = (1, 1, 2, 2, 4, 4)
-    conf.net_enc_channel_mult = (1, 1, 2, 2, 4, 4, 4)
-    conf.eval_every_samples = 10_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.total_samples = 200_000_000
-    conf.batch_size = 64
-    conf.make_model_conf()
-    conf.name = 'ffhq256_autoenc_eco'
-    return conf
-
-
-def ffhq128_ddpm_72M():
-    conf = ffhq128_ddpm()
-    conf.total_samples = 72_000_000
-    conf.name = 'ffhq128_ddpm_72M'
-    return conf
-
-
-def ffhq128_autoenc_72M():
-    conf = ffhq128_autoenc_base()
-    conf.total_samples = 72_000_000
-    conf.name = 'ffhq128_autoenc_72M'
-    return conf
-
-
-def ffhq128_ddpm_130M():
-    conf = ffhq128_ddpm()
-    conf.total_samples = 130_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.name = 'ffhq128_ddpm_130M'
-    return conf
-
-
-def ffhq128_autoenc_130M():
-    conf = ffhq128_autoenc_base()
-    conf.total_samples = 130_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.name = 'ffhq128_autoenc_130M'
-    return conf
-
-
-def horse128_ddpm():
-    conf = ffhq128_ddpm()
-    conf.data_name = 'horse256'
-    conf.total_samples = 130_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.name = 'horse128_ddpm'
-    return conf
-
-
-def horse128_autoenc():
-    conf = ffhq128_autoenc_base()
-    conf.data_name = 'horse256'
-    conf.total_samples = 130_000_000
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.name = 'horse128_autoenc'
-    return conf
-
-
-def bedroom128_ddpm():
-    conf = ffhq128_ddpm()
-    conf.data_name = 'bedroom256'
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.total_samples = 120_000_000
-    conf.name = 'bedroom128_ddpm'
-    return conf
-
-
-def bedroom128_autoenc():
-    conf = ffhq128_autoenc_base()
-    conf.data_name = 'bedroom256'
-    conf.eval_ema_every_samples = 10_000_000
-    conf.eval_every_samples = 10_000_000
-    conf.total_samples = 120_000_000
-    conf.name = 'bedroom128_autoenc'
-    return conf
-
-
-def pretrain_celeba64d2c_72M():
-    conf = celeba64d2c_autoenc()
-    conf.pretrain = PretrainConfig(
-        name='72M',
-        path=f'checkpoints/{celeba64d2c_autoenc().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{celeba64d2c_autoenc().name}/latent.pkl'
-    return conf
-
-
-def pretrain_ffhq128_autoenc72M():
-    conf = ffhq128_autoenc_base()
-    conf.postfix = ''
-    conf.pretrain = PretrainConfig(
-        name='72M',
-        path=f'checkpoints/{ffhq128_autoenc_72M().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{ffhq128_autoenc_72M().name}/latent.pkl'
-    return conf
-
-
-def pretrain_ffhq128_autoenc130M():
-    conf = ffhq128_autoenc_base()
-    conf.pretrain = PretrainConfig(
-        name='130M',
-        path=f'checkpoints/{ffhq128_autoenc_130M().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{ffhq128_autoenc_130M().name}/latent.pkl'
-    return conf
-
-
-def pretrain_ffhq256_autoenc():
-    conf = ffhq256_autoenc()
-    conf.pretrain = PretrainConfig(
-        name='90M',
-        path=f'checkpoints/{ffhq256_autoenc().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{ffhq256_autoenc().name}/latent.pkl'
-    return conf
-
-
-def pretrain_horse128():
-    conf = horse128_autoenc()
-    conf.pretrain = PretrainConfig(
-        name='82M',
-        path=f'checkpoints/{horse128_autoenc().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{horse128_autoenc().name}/latent.pkl'
-    return conf
-
-
-def pretrain_bedroom128():
-    conf = bedroom128_autoenc()
-    conf.pretrain = PretrainConfig(
-        name='120M',
-        path=f'checkpoints/{bedroom128_autoenc().name}/last.ckpt',
-    )
-    conf.latent_infer_path = f'checkpoints/{bedroom128_autoenc().name}/latent.pkl'
-    return conf