Upload 49 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/id02548.0pAkJZmlFqc.00001_id04570.0YMGn6BI9rg.00001.gif filter=lfs diff=lfs merge=lfs -text
+assets/website_gif_v2.gif filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import subprocess
+import os
+import requests
+
+def process_video(audio_file, video_file):
+    # Define file paths
+    audio_path = audio_file.name
+    video_path = video_file.name
+    out_path = "output_video.mp4"
+
+    # Save uploaded files
+    audio_file.save(audio_path)
+    video_file.save(video_path)
+
+    # Define command flags
+    sample_mode = "cross"  # or "reconstruction"
+    generate_from_filelist = 0
+    model_path = "checkpoints/checkpoint.pt"
+    pads = "0,0,0,0"
+
+    if sample_mode == "reconstruction":
+        sample_input_flags = "--sampling_input_type=first_frame --sampling_ref_type=first_frame"
+    elif sample_mode == "cross":
+        sample_input_flags = "--sampling_input_type=gt --sampling_ref_type=gt"
+    else:
+        return "Error: sample_mode can only be \"cross\" or \"reconstruction\""
+
+    MODEL_FLAGS = "--attention_resolutions 32,16,8 --class_cond False --learn_sigma True --num_channels 128 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm False"
+    DIFFUSION_FLAGS = "--predict_xstart False --diffusion_steps 1000 --noise_schedule linear --rescale_timesteps False"
+    SAMPLE_FLAGS = f"--sampling_seed=7 {sample_input_flags} --timestep_respacing ddim25 --use_ddim True --model_path={model_path}"
+    DATA_FLAGS = "--nframes 5 --nrefer 1 --image_size 128 --sampling_batch_size=32"
+    TFG_FLAGS = "--face_hide_percentage 0.5 --use_ref=True --use_audio=True --audio_as_style=True"
+    GEN_FLAGS = f"--generate_from_filelist {generate_from_filelist} --video_path={video_path} --audio_path={audio_path} --out_path={out_path} --save_orig=False --face_det_batch_size 16 --pads {pads} --is_voxceleb2=False"
+
+    # Combine all flags into one command
+    command = f"python your_model_script.py {MODEL_FLAGS} {DIFFUSION_FLAGS} {SAMPLE_FLAGS} {DATA_FLAGS} {TFG_FLAGS} {GEN_FLAGS}"
+
+    # Execute the command
+    try:
+        subprocess.run(command, shell=True, check=True)
+        return out_path
+    except subprocess.CalledProcessError as e:
+        return f"Error processing video: {e}"
+
+    # Clean up the files after processing
+    os.remove(audio_path)
+    os.remove(video_path)
+
+    # Delete output video after sending to the user
+    os.remove(out_path)
+
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=process_video,
+    inputs=[
+        gr.inputs.Audio(label="Input Audio", type="file"),
+        gr.inputs.Video(label="Input Video", type="file")
+    ],
+    outputs=gr.outputs.Video(label="Processed Video"),
+    title="Audio-Video Processing",
+    description="Upload an audio file and a video file to process the video based on the audio input."
+)
+
+# Launch the interface
+iface.launch()

audio/audio.py

@@ -0,0 +1,136 @@
+import librosa
+import librosa.filters
+import numpy as np
+# import tensorflow as tf
+from scipy import signal
+from scipy.io import wavfile
+from .hparams import hparams as hp
+
+def load_wav(path, sr):
+    return librosa.core.load(path, sr=sr)[0]
+
+def save_wav(wav, path, sr):
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    #proposed by @dsmiller
+    wavfile.write(path, sr, wav.astype(np.int16))
+
+def save_wavenet_wav(wav, path, sr):
+    librosa.output.write_wav(path, wav, sr=sr)
+
+def preemphasis(wav, k, preemphasize=True):
+    if preemphasize:
+        return signal.lfilter([1, -k], [1], wav)
+    return wav
+
+def inv_preemphasis(wav, k, inv_preemphasize=True):
+    if inv_preemphasize:
+        return signal.lfilter([1], [1, -k], wav)
+    return wav
+
+def get_hop_size():
+    hop_size = hp.hop_size
+    if hop_size is None:
+        assert hp.frame_shift_ms is not None
+        hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
+    return hop_size
+
+def linearspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(np.abs(D)) - hp.ref_level_db
+    
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+def melspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
+    
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+def _lws_processor():
+    import lws
+    return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
+
+def _stft(y):
+    if hp.use_lws:
+        return _lws_processor(hp).stft(y).T
+    else:
+        return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size)
+
+##########################################################
+#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
+def num_frames(length, fsize, fshift):
+    """Compute number of time frames of spectrogram
+    """
+    pad = (fsize - fshift)
+    if length % fshift == 0:
+        M = (length + pad * 2 - fsize) // fshift + 1
+    else:
+        M = (length + pad * 2 - fsize) // fshift + 2
+    return M
+
+
+def pad_lr(x, fsize, fshift):
+    """Compute left and right padding
+    """
+    M = num_frames(len(x), fsize, fshift)
+    pad = (fsize - fshift)
+    T = len(x) + 2 * pad
+    r = (M - 1) * fshift + fsize - T
+    return pad, pad + r
+##########################################################
+#Librosa correct padding
+def librosa_pad_lr(x, fsize, fshift):
+    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
+
+# Conversions
+_mel_basis = None
+
+def _linear_to_mel(spectogram):
+    global _mel_basis
+    if _mel_basis is None:
+        _mel_basis = _build_mel_basis()
+    return np.dot(_mel_basis, spectogram)
+
+def _build_mel_basis():
+    assert hp.fmax <= hp.sample_rate // 2
+    return librosa.filters.mel(sr=hp.sample_rate, n_fft=hp.n_fft, n_mels=hp.num_mels,
+                               fmin=hp.fmin, fmax=hp.fmax)
+
+def _amp_to_db(x):
+    min_level = np.exp(hp.min_level_db / 20 * np.log(10))
+    return 20 * np.log10(np.maximum(min_level, x))
+
+def _db_to_amp(x):
+    return np.power(10.0, (x) * 0.05)
+
+def _normalize(S):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value,
+                           -hp.max_abs_value, hp.max_abs_value)
+        else:
+            return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value)
+    
+    assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
+    if hp.symmetric_mels:
+        return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value
+    else:
+        return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
+
+def _denormalize(D):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return (((np.clip(D, -hp.max_abs_value,
+                              hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value))
+                    + hp.min_level_db)
+        else:
+            return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
+    
+    if hp.symmetric_mels:
+        return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db)
+    else:
+        return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)

audio/hparams.py

@@ -0,0 +1,66 @@
+from glob import glob
+import os
+
+
+class HParams:
+	def __init__(self, **kwargs):
+		self.data = {}
+
+		for key, value in kwargs.items():
+			self.data[key] = value
+
+	def __getattr__(self, key):
+		if key not in self.data:
+			raise AttributeError("'HParams' object has no attribute %s" % key)
+		return self.data[key]
+
+	def set_hparam(self, key, value):
+		self.data[key] = value
+
+
+# Default hyperparameters
+hparams = HParams(
+	num_mels=80,  # Number of mel-spectrogram channels and local conditioning dimensionality
+	#  network
+	rescale=True,  # Whether to rescale audio prior to preprocessing
+	rescaling_max=0.9,  # Rescaling value
+	
+	# Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
+	# It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
+	# Does not work if n_ffit is not multiple of hop_size!!
+	use_lws=False,
+	
+	n_fft=800,  # Extra window size is filled with 0 paddings to match this parameter
+	hop_size=200,  # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
+	win_size=800,  # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
+	sample_rate=16000,  # 16000Hz (corresponding to librispeech) (sox --i <filename>)
+	
+	frame_shift_ms=None,  # Can replace hop_size parameter. (Recommended: 12.5)
+	
+	# Mel and Linear spectrograms normalization/scaling and clipping
+	signal_normalization=True,
+	# Whether to normalize mel spectrograms to some predefined range (following below parameters)
+	allow_clipping_in_normalization=True,  # Only relevant if mel_normalization = True
+	symmetric_mels=True,
+	# Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2, 
+	# faster and cleaner convergence)
+	max_abs_value=4.,
+	# max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not 
+	# be too big to avoid gradient explosion, 
+	# not too small for fast convergence)
+	# Contribution by @begeekmyfriend
+	# Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude 
+	# levels. Also allows for better G&L phase reconstruction)
+	preemphasize=True,  # whether to apply filter
+	preemphasis=0.97,  # filter coefficient.
+	
+	# Limits
+	min_level_db=-100,
+	ref_level_db=20,
+	fmin=55,
+	# Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To 
+	# test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
+	fmax=7600,  # To be increased/reduced depending on data.
+
+)
+

checkpoints/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c71166482d2b893f2f77450563a1bb31d805f3048c7213b974fd9201e9aa4b3
+size 406815527

dataset/LRW/lrw_fullpath.py

@@ -0,0 +1,25 @@
+'''Converts the LRW video names in filelists to LRW relative paths and dumps them unto new filelists'''
+import os
+filelist = "../filelists/lrw_cross.txt" 
+
+filelist_split_path = filelist.replace(".txt","_relative_path.txt")
+with open(filelist, 'r') as f:
+        lines = f.readlines()
+with open(filelist_split_path, 'w') as f:
+    for i in range(len(lines)):
+        audio_name, video_name=lines[i].split(' ')
+        audio_word = audio_name.split('_')[0]
+        video_word = video_name.split('_')[0]
+        f.write(os.path.join(audio_word,'test',audio_name)+' '+os.path.join(video_word,'test',video_name))
+
+filelist = "../filelists/lrw_reconstruction.txt" 
+
+filelist_split_path = filelist.replace(".txt","_relative_path.txt")
+with open(filelist, 'r') as f:
+        lines = f.readlines()
+with open(filelist_split_path, 'w') as f:
+    for i in range(len(lines)):
+        audio_name, video_name=lines[i].split(' ')
+        audio_word = audio_name.split('_')[0]
+        video_word = video_name.split('_')[0]
+        f.write(os.path.join(audio_word,'test',audio_name)+' '+os.path.join(video_word,'test',video_name))

dataset/filelists/voxceleb2_test_n_500_reconstruction.txt

@@ -0,0 +1,500 @@
+id09017/SjCgiXBHfNU/00111 id09017/SjCgiXBHfNU/00111
+id05055/HobsYUHmgr0/00138 id05055/HobsYUHmgr0/00138
+id01567/M47d5UckOV8/00099 id01567/M47d5UckOV8/00099
+id01228/SH3eBLMsRwY/00211 id01228/SH3eBLMsRwY/00211
+id07312/m1VY1sC_P_o/00093 id07312/m1VY1sC_P_o/00093
+id08696/dJH9aBSs1nE/00370 id08696/dJH9aBSs1nE/00370
+id07312/9PCY4xwxgcE/00006 id07312/9PCY4xwxgcE/00006
+id07494/1A8ZDo11tzY/00006 id07494/1A8ZDo11tzY/00006
+id00061/thHLZ8tDJ-M/00276 id00061/thHLZ8tDJ-M/00276
+id03862/YQMdTzyG-P8/00297 id03862/YQMdTzyG-P8/00297
+id04570/zsnG6eKzOGE/00406 id04570/zsnG6eKzOGE/00406
+id07414/1m1C-CdhmZ0/00016 id07414/1m1C-CdhmZ0/00016
+id01509/2uKpHd-euIo/00038 id01509/2uKpHd-euIo/00038
+id04276/qHXwXqxL0mk/00401 id04276/qHXwXqxL0mk/00401
+id04366/x-VQ6z2QC4w/00252 id04366/x-VQ6z2QC4w/00252
+id07620/um6sY627GaE/00475 id07620/um6sY627GaE/00475
+id01000/RvjbLfo3XDM/00052 id01000/RvjbLfo3XDM/00052
+id07868/fnWDbUI_Zbg/00289 id07868/fnWDbUI_Zbg/00289
+id01333/cymDCPEhalE/00351 id01333/cymDCPEhalE/00351
+id02317/Mv16h1Bx7HE/00241 id02317/Mv16h1Bx7HE/00241
+id02317/Vi4k3cuwfgc/00342 id02317/Vi4k3cuwfgc/00342
+id01000/eGeGHhuOJJ0/00077 id01000/eGeGHhuOJJ0/00077
+id03980/zaDLb12pDBQ/00130 id03980/zaDLb12pDBQ/00130
+id05124/c-Pa7b81coQ/00354 id05124/c-Pa7b81coQ/00354
+id04478/nhLuGj0vGb8/00234 id04478/nhLuGj0vGb8/00234
+id01541/3su8tn9nwi4/00007 id01541/3su8tn9nwi4/00007
+id06484/cmIyVotzXiE/00125 id06484/cmIyVotzXiE/00125
+id06209/oxofNHGCj7s/00139 id06209/oxofNHGCj7s/00139
+id02181/rxX3t2rzLbg/00146 id02181/rxX3t2rzLbg/00146
+id02286/YL75-u9XYUM/00105 id02286/YL75-u9XYUM/00105
+id04276/v9mSslwD0Kg/00470 id04276/v9mSslwD0Kg/00470
+id07802/6qBSFfV_Mig/00042 id07802/6qBSFfV_Mig/00042
+id04295/DtC2X1KG8TE/00057 id04295/DtC2X1KG8TE/00057
+id00866/shG_183xFlw/00243 id00866/shG_183xFlw/00243
+id03862/2nagLhV_Yvw/00012 id03862/2nagLhV_Yvw/00012
+id04119/Yndoy1jgHWs/00042 id04119/Yndoy1jgHWs/00042
+id04295/mTCDT_Fv5So/00203 id04295/mTCDT_Fv5So/00203
+id08456/o-5hKwhGqac/00354 id08456/o-5hKwhGqac/00354
+id07494/tv6GJkx_Wy4/00331 id07494/tv6GJkx_Wy4/00331
+id04295/mClPHVzTCLI/00196 id04295/mClPHVzTCLI/00196
+id04478/81Tb6kjlNIk/00019 id04478/81Tb6kjlNIk/00019
+id00812/NeNXGI8mox8/00158 id00812/NeNXGI8mox8/00158
+id04536/sfldoEPrFPI/00438 id04536/sfldoEPrFPI/00438
+id07620/aJVbccKJwEw/00327 id07620/aJVbccKJwEw/00327
+id02286/4LAIxvdvguc/00001 id02286/4LAIxvdvguc/00001
+id07802/BfQUBDw7TiM/00080 id07802/BfQUBDw7TiM/00080
+id01066/65k0p7fUBVI/00026 id01066/65k0p7fUBVI/00026
+id03862/w97YzyPYm1k/00460 id03862/w97YzyPYm1k/00460
+id05816/njFBkJSpUrY/00414 id05816/njFBkJSpUrY/00414
+id05124/_Oxp6absIhY/00341 id05124/_Oxp6absIhY/00341
+id07663/mUw-kxAavdM/00192 id07663/mUw-kxAavdM/00192
+id05999/Ls440srvfR4/00127 id05999/Ls440srvfR4/00127
+id02548/Hmlw5PIf64o/00098 id02548/Hmlw5PIf64o/00098
+id04276/Pbo_nlcZ0Lc/00190 id04276/Pbo_nlcZ0Lc/00190
+id07802/FhKML4dLE60/00115 id07802/FhKML4dLE60/00115
+id07621/1L2IUy6gqaM/00012 id07621/1L2IUy6gqaM/00012
+id05654/veGIQ7p2ZSk/00130 id05654/veGIQ7p2ZSk/00130
+id04094/0z1JYPKGBI8/00007 id04094/0z1JYPKGBI8/00007
+id02576/wWUREnOwYo0/00136 id02576/wWUREnOwYo0/00136
+id09017/PLNK1g5w4FY/00099 id09017/PLNK1g5w4FY/00099
+id06484/USbx34RUkVI/00096 id06484/USbx34RUkVI/00096
+id03030/FXbzdRO7t98/00101 id03030/FXbzdRO7t98/00101
+id02057/VCXnx-ozS8c/00263 id02057/VCXnx-ozS8c/00263
+id02542/JUodrwt9ucI/00033 id02542/JUodrwt9ucI/00033
+id03030/DM_Z5D2fkRA/00068 id03030/DM_Z5D2fkRA/00068
+id08552/irj3SqKAe0c/00196 id08552/irj3SqKAe0c/00196
+id03030/YxBoufnVIMw/00177 id03030/YxBoufnVIMw/00177
+id07868/Eaf-dgA59Gs/00061 id07868/Eaf-dgA59Gs/00061
+id08456/6xVSlQDr7-w/00031 id08456/6xVSlQDr7-w/00031
+id06811/OYFkt_n18hg/00128 id06811/OYFkt_n18hg/00128
+id00817/tCnW5E8cMow/00383 id00817/tCnW5E8cMow/00383
+id02542/fXQbNcIbcek/00053 id02542/fXQbNcIbcek/00053
+id01567/oi2g17EF55s/00377 id01567/oi2g17EF55s/00377
+id04366/HsG3OGE22DY/00117 id04366/HsG3OGE22DY/00117
+id01509/1y0aWmgYDtw/00006 id01509/1y0aWmgYDtw/00006
+id04295/pYfyopS672Y/00213 id04295/pYfyopS672Y/00213
+id01989/6JfW9CPAoGY/00006 id01989/6JfW9CPAoGY/00006
+id04366/tbcKV-IjZdI/00243 id04366/tbcKV-IjZdI/00243
+id01298/UY0fkYSUFrY/00208 id01298/UY0fkYSUFrY/00208
+id00817/GAs8WnyFKJM/00120 id00817/GAs8WnyFKJM/00120
+id06484/TCp2-XVatIE/00079 id06484/TCp2-XVatIE/00079
+id08374/Kf9N5AWprG8/00150 id08374/Kf9N5AWprG8/00150
+id01822/QDWgjZqOkvM/00065 id01822/QDWgjZqOkvM/00065
+id03030/pTz652Dx_6w/00230 id03030/pTz652Dx_6w/00230
+id01460/chrI43l2Nuw/00201 id01460/chrI43l2Nuw/00201
+id08374/85f-qB_KJP8/00041 id08374/85f-qB_KJP8/00041
+id07961/PoSkUxZ4ags/00172 id07961/PoSkUxZ4ags/00172
+id01437/uFPYqotT7tU/00233 id01437/uFPYqotT7tU/00233
+id07621/Aan8MoozxII/00095 id07621/Aan8MoozxII/00095
+id08456/fWTULQWYVoA/00250 id08456/fWTULQWYVoA/00250
+id05055/da7Z8oWhFPY/00351 id05055/da7Z8oWhFPY/00351
+id02181/hIvctbfcBx8/00106 id02181/hIvctbfcBx8/00106
+id01541/dEmuPb4A7do/00184 id01541/dEmuPb4A7do/00184
+id00419/a3Y7pQzcn40/00305 id00419/a3Y7pQzcn40/00305
+id07354/dsDxN33xvL0/00262 id07354/dsDxN33xvL0/00262
+id04478/MZh3AEgJ9pc/00092 id04478/MZh3AEgJ9pc/00092
+id05124/UBUFmICrT-I/00281 id05124/UBUFmICrT-I/00281
+id03127/SmGJu-t24hY/00195 id03127/SmGJu-t24hY/00195
+id02465/coOp_DnsmEI/00150 id02465/coOp_DnsmEI/00150
+id01618/qrOl1aaXBH0/00187 id01618/qrOl1aaXBH0/00187
+id03969/WZVnB-m0X9g/00038 id03969/WZVnB-m0X9g/00038
+id05202/s0m_4-SCn44/00186 id05202/s0m_4-SCn44/00186
+id04657/SYVkfHq-pro/00172 id04657/SYVkfHq-pro/00172
+id05176/p2IOP5_s_LM/00093 id05176/p2IOP5_s_LM/00093
+id04950/XJS6SLQuCNM/00169 id04950/XJS6SLQuCNM/00169
+id02019/anSrwA_9RPE/00152 id02019/anSrwA_9RPE/00152
+id04570/Q-faEy1VXxQ/00140 id04570/Q-faEy1VXxQ/00140
+id07621/bMvG2mQMZZw/00303 id07621/bMvG2mQMZZw/00303
+id06811/vC3yQiWuuOI/00354 id06811/vC3yQiWuuOI/00354
+id03839/aWMP8xzq2BE/00292 id03839/aWMP8xzq2BE/00292
+id04094/j1ajUkR6_Q4/00326 id04094/j1ajUkR6_Q4/00326
+id08149/o0Zdr9Jla7U/00047 id08149/o0Zdr9Jla7U/00047
+id00017/hcr4tT9y3xs/00117 id00017/hcr4tT9y3xs/00117
+id04950/Cu4jGRmYa4c/00064 id04950/Cu4jGRmYa4c/00064
+id01567/TMozlhoPMfI/00223 id01567/TMozlhoPMfI/00223
+id08374/QltFme-lqeI/00226 id08374/QltFme-lqeI/00226
+id06816/tHor4VN8090/00259 id06816/tHor4VN8090/00259
+id07494/xQ0YMPe-9u8/00413 id07494/xQ0YMPe-9u8/00413
+id08374/FwR1K1rL3QI/00110 id08374/FwR1K1rL3QI/00110
+id06692/Hlahj5abifM/00257 id06692/Hlahj5abifM/00257
+id00419/J2LscHjRX7Q/00154 id00419/J2LscHjRX7Q/00154
+id02057/CI5-q_qTR5I/00112 id02057/CI5-q_qTR5I/00112
+id03862/7IccaH4HXRs/00069 id03862/7IccaH4HXRs/00069
+id04656/ar3rKrkbjqI/00257 id04656/ar3rKrkbjqI/00257
+id07494/XMEIdqio6ic/00184 id07494/XMEIdqio6ic/00184
+id04657/dn4XY5c6mEw/00265 id04657/dn4XY5c6mEw/00265
+id04570/SFKt669qIqs/00156 id04570/SFKt669qIqs/00156
+id01541/sMDYdAB0MPs/00306 id01541/sMDYdAB0MPs/00306
+id08456/F2O-frqyr9c/00101 id08456/F2O-frqyr9c/00101
+id08701/_Ysb9mVibbk/00253 id08701/_Ysb9mVibbk/00253
+id01333/e4FoER8nqx0/00365 id01333/e4FoER8nqx0/00365
+id05124/F0Xpd6OoiDY/00161 id05124/F0Xpd6OoiDY/00161
+id01593/AVmZf6Kl1So/00071 id01593/AVmZf6Kl1So/00071
+id01567/fOlxxDqdrgc/00299 id01567/fOlxxDqdrgc/00299
+id06484/2KVWoftPf2o/00001 id06484/2KVWoftPf2o/00001
+id01224/g4jVqkEm1Gs/00274 id01224/g4jVqkEm1Gs/00274
+id02445/ZX_6RMrTEP0/00066 id02445/ZX_6RMrTEP0/00066
+id04656/5TR-W77XgF4/00032 id04656/5TR-W77XgF4/00032
+id01618/F_ExF9xDajc/00060 id01618/F_ExF9xDajc/00060
+id08392/gPX4IC53KwI/00355 id08392/gPX4IC53KwI/00355
+id00866/pNbDtfW1JW4/00221 id00866/pNbDtfW1JW4/00221
+id00812/b3dBqOtzsx0/00276 id00812/b3dBqOtzsx0/00276
+id08701/61Al05HARgA/00001 id08701/61Al05HARgA/00001
+id07663/FFo4JwVXeUM/00119 id07663/FFo4JwVXeUM/00119
+id02057/22zJ50ky7CQ/00013 id02057/22zJ50ky7CQ/00013
+id05055/2onVoeSgouI/00028 id05055/2onVoeSgouI/00028
+id04006/zvUZFL0NyhM/00260 id04006/zvUZFL0NyhM/00260
+id04950/EpOnsaBin0A/00077 id04950/EpOnsaBin0A/00077
+id05015/RhBpC9Fc7a4/00154 id05015/RhBpC9Fc7a4/00154
+id04656/Z_JFBDW9eZE/00251 id04656/Z_JFBDW9eZE/00251
+id01509/2sb83ZBlbJg/00034 id01509/2sb83ZBlbJg/00034
+id04030/JbcD0P6KGe0/00036 id04030/JbcD0P6KGe0/00036
+id02542/cwgUjse_REU/00040 id02542/cwgUjse_REU/00040
+id07620/xFc9X6EXtRM/00478 id07620/xFc9X6EXtRM/00478
+id07354/Qrg89rvtZ1k/00217 id07354/Qrg89rvtZ1k/00217
+id03839/wSQMEZMxxx4/00461 id03839/wSQMEZMxxx4/00461
+id03127/iWeklsXc0H8/00268 id03127/iWeklsXc0H8/00268
+id07663/54qlJ2HZ08s/00096 id07663/54qlJ2HZ08s/00096
+id07961/Orp8s5aHYc8/00158 id07961/Orp8s5aHYc8/00158
+id03347/y_F4aAkN0d8/00417 id03347/y_F4aAkN0d8/00417
+id06913/KNDyf594xQg/00056 id06913/KNDyf594xQg/00056
+id04366/DIgAc22fq9c/00080 id04366/DIgAc22fq9c/00080
+id07396/uJPtbxlXi2c/00187 id07396/uJPtbxlXi2c/00187
+id07868/gVspdH-U2XE/00290 id07868/gVspdH-U2XE/00290
+id05594/u7qCFBP1nH4/00184 id05594/u7qCFBP1nH4/00184
+id01541/mDoT5mpo_2c/00241 id01541/mDoT5mpo_2c/00241
+id07354/0y9b8qlM170/00011 id07354/0y9b8qlM170/00011
+id01460/DnnphhTlRPE/00075 id01460/DnnphhTlRPE/00075
+id02548/1CNhmMmirfA/00009 id02548/1CNhmMmirfA/00009
+id03127/k8z6DxdyF9w/00291 id03127/k8z6DxdyF9w/00291
+id01437/zLRJ_8_M5Wg/00263 id01437/zLRJ_8_M5Wg/00263
+id02576/WnbNQuJzErQ/00086 id02576/WnbNQuJzErQ/00086
+id01333/M0UD9g1x18c/00128 id01333/M0UD9g1x18c/00128
+id04295/1fSjOItVYVg/00001 id04295/1fSjOItVYVg/00001
+id08456/8tt1LbCoU0E/00054 id08456/8tt1LbCoU0E/00054
+id07494/r-ToqH_EJNs/00318 id07494/r-ToqH_EJNs/00318
+id06816/XBKj9XWlZCw/00123 id06816/XBKj9XWlZCw/00123
+id03030/haoNit7a4W0/00201 id03030/haoNit7a4W0/00201
+id03839/aeObhOJLQzQ/00293 id03839/aeObhOJLQzQ/00293
+id07868/COb1gFHXsBQ/00059 id07868/COb1gFHXsBQ/00059
+id01224/eYWcMCsgkLY/00255 id01224/eYWcMCsgkLY/00255
+id04006/K5ueXBlS6rc/00049 id04006/K5ueXBlS6rc/00049
+id07620/G5-1CUbaz0c/00107 id07620/G5-1CUbaz0c/00107
+id06104/cj0TAnwndoc/00230 id06104/cj0TAnwndoc/00230
+id00061/STX1ycPt8fU/00076 id00061/STX1ycPt8fU/00076
+id04478/wMbobxEQ7j8/00336 id04478/wMbobxEQ7j8/00336
+id01106/7X_xtnJhEc0/00031 id01106/7X_xtnJhEc0/00031
+id08374/zaYzRbE_2C8/00494 id08374/zaYzRbE_2C8/00494
+id04276/MgOqCfwKE70/00173 id04276/MgOqCfwKE70/00173
+id03127/Lgd5qn2-kMo/00079 id03127/Lgd5qn2-kMo/00079
+id00154/xH3Pp_5yxOk/00153 id00154/xH3Pp_5yxOk/00153
+id04030/7mXUMuo5_NE/00001 id04030/7mXUMuo5_NE/00001
+id02542/p7bvjcLbZm4/00097 id02542/p7bvjcLbZm4/00097
+id04232/T7dROCqmwNQ/00235 id04232/T7dROCqmwNQ/00235
+id02548/KrXU-_jrtxY/00147 id02548/KrXU-_jrtxY/00147
+id01567/SZyTC5dxJOY/00219 id01567/SZyTC5dxJOY/00219
+id03524/2DD4Np7SaWw/00007 id03524/2DD4Np7SaWw/00007
+id04094/DRq5F2261Ko/00072 id04094/DRq5F2261Ko/00072
+id07802/HrpJg06dowY/00152 id07802/HrpJg06dowY/00152
+id06816/pBt-DxsTFc8/00231 id06816/pBt-DxsTFc8/00231
+id00154/2pSNL5YdcoQ/00002 id00154/2pSNL5YdcoQ/00002
+id01541/C29fUBtimOE/00038 id01541/C29fUBtimOE/00038
+id06310/b6qPjJ0isPI/00155 id06310/b6qPjJ0isPI/00155
+id05714/wFGNufaMbDY/00025 id05714/wFGNufaMbDY/00025
+id03980/m-8Ffv2RqYs/00092 id03980/m-8Ffv2RqYs/00092
+id01437/uXAe0vbNWeo/00238 id01437/uXAe0vbNWeo/00238
+id04232/tPZ-zVT67gs/00479 id04232/tPZ-zVT67gs/00479
+id06811/ImzUwwYU6SQ/00067 id06811/ImzUwwYU6SQ/00067
+id05459/wq3Z0I944wU/00436 id05459/wq3Z0I944wU/00436
+id03969/Evoldg-U2_c/00024 id03969/Evoldg-U2_c/00024
+id08548/BSChFozahbU/00019 id08548/BSChFozahbU/00019
+id04950/PQEAck-3wcA/00134 id04950/PQEAck-3wcA/00134
+id04295/G4YnExZSzlM/00066 id04295/G4YnExZSzlM/00066
+id05176/mc7rFp2B1j0/00092 id05176/mc7rFp2B1j0/00092
+id00812/1Xfgvdu7oDo/00001 id00812/1Xfgvdu7oDo/00001
+id05459/UPSPGawaVsg/00233 id05459/UPSPGawaVsg/00233
+id04656/7nG3rOv0oBw/00050 id04656/7nG3rOv0oBw/00050
+id02548/nvYBpt14BrQ/00309 id02548/nvYBpt14BrQ/00309
+id02317/A3AvljK8Upk/00102 id02317/A3AvljK8Upk/00102
+id04478/qLNvRwMkhik/00242 id04478/qLNvRwMkhik/00242
+id01228/lCDMC8JvKyU/00295 id01228/lCDMC8JvKyU/00295
+id03041/5CfnYwQCW48/00001 id03041/5CfnYwQCW48/00001
+id04950/LnsriCjCIV4/00116 id04950/LnsriCjCIV4/00116
+id04094/plxNYSFgDTM/00384 id04094/plxNYSFgDTM/00384
+id01460/30_QmGw7lmE/00030 id01460/30_QmGw7lmE/00030
+id04366/6rX7hCNSjaw/00056 id04366/6rX7hCNSjaw/00056
+id01041/m-xolqIq8p4/00370 id01041/m-xolqIq8p4/00370
+id04950/BG4CCg2RiuQ/00052 id04950/BG4CCg2RiuQ/00052
+id01989/7g0A7pF94r0/00018 id01989/7g0A7pF94r0/00018
+id03382/b_NJ2Xz3G4Y/00030 id03382/b_NJ2Xz3G4Y/00030
+id00812/IteHRVKyzaE/00138 id00812/IteHRVKyzaE/00138
+id00061/bdkqfVtDZVY/00121 id00061/bdkqfVtDZVY/00121
+id03839/YkYIh4cYwwg/00275 id03839/YkYIh4cYwwg/00275
+id07354/wyTuCRGjUIQ/00477 id07354/wyTuCRGjUIQ/00477
+id02057/TddnW2TaXrc/00246 id02057/TddnW2TaXrc/00246
+id01989/gHVHtKTQBsw/00128 id01989/gHVHtKTQBsw/00128
+id08374/bXlUHb5hxxA/00266 id08374/bXlUHb5hxxA/00266
+id03862/TE2zQc8_W-g/00252 id03862/TE2zQc8_W-g/00252
+id08696/86-k8TuowAE/00033 id08696/86-k8TuowAE/00033
+id05176/K8yZYHg_4ro/00050 id05176/K8yZYHg_4ro/00050
+id04253/SKsPkHMGHYY/00240 id04253/SKsPkHMGHYY/00240
+id07874/2KK4ozkjaEE/00002 id07874/2KK4ozkjaEE/00002
+id08392/g-SJYYaaLgE/00352 id08392/g-SJYYaaLgE/00352
+id02542/glhCf1hwJhE/00065 id02542/glhCf1hwJhE/00065
+id00817/FsL-bTbDTyw/00112 id00817/FsL-bTbDTyw/00112
+id04862/IuXPj9VhUVA/00100 id04862/IuXPj9VhUVA/00100
+id06811/f9-8d3lNNcw/00237 id06811/f9-8d3lNNcw/00237
+id04094/JUYMzfVp8zI/00113 id04094/JUYMzfVp8zI/00113
+id03347/r-xJUB0A4ok/00346 id03347/r-xJUB0A4ok/00346
+id07868/MNibTv_ODQ8/00148 id07868/MNibTv_ODQ8/00148
+id08392/3e5zvNaT-eU/00020 id08392/3e5zvNaT-eU/00020
+id04295/bKMKvAr440A/00141 id04295/bKMKvAr440A/00141
+id04295/l62YPD0ZkZI/00185 id04295/l62YPD0ZkZI/00185
+id07312/RO9DsspwXiE/00047 id07312/RO9DsspwXiE/00047
+id03030/rmFsUV5ICKk/00267 id03030/rmFsUV5ICKk/00267
+id03677/nVWTTopGQdU/00181 id03677/nVWTTopGQdU/00181
+id00866/xQ1Yy0kjvjA/00256 id00866/xQ1Yy0kjvjA/00256
+id01333/fRnqtJR0rws/00371 id01333/fRnqtJR0rws/00371
+id05055/AZoIKG33E8s/00115 id05055/AZoIKG33E8s/00115
+id01822/_CkfCmQXII8/00098 id01822/_CkfCmQXII8/00098
+id01593/_gyaAyVi6SA/00344 id01593/_gyaAyVi6SA/00344
+id04295/DS3RDwf2xI8/00049 id04295/DS3RDwf2xI8/00049
+id00812/EjO-VORTv_o/00098 id00812/EjO-VORTv_o/00098
+id04657/WdJ_DuU0ack/00236 id04657/WdJ_DuU0ack/00236
+id04232/AB9fk1MH2rA/00035 id04232/AB9fk1MH2rA/00035
+id00419/chfgCUm9-Mg/00364 id00419/chfgCUm9-Mg/00364
+id02577/Az0BGrX_TwI/00021 id02577/Az0BGrX_TwI/00021
+id01437/hyj4OYm0cvA/00195 id01437/hyj4OYm0cvA/00195
+id01593/tLFWX-IdAwI/00431 id01593/tLFWX-IdAwI/00431
+id04536/MNDmkEXRS7s/00312 id04536/MNDmkEXRS7s/00312
+id03789/7qhkM8qY3Fw/00077 id03789/7qhkM8qY3Fw/00077
+id01593/neAk6K8BvTA/00397 id01593/neAk6K8BvTA/00397
+id06484/jTHSVo6NvS4/00151 id06484/jTHSVo6NvS4/00151
+id07414/cAudd_5Yv2I/00256 id07414/cAudd_5Yv2I/00256
+id00866/ADzqaRZtJNA/00087 id00866/ADzqaRZtJNA/00087
+id06484/ZySpn0Aj09k/00108 id06484/ZySpn0Aj09k/00108
+id07312/ZHBjHQENqW8/00053 id07312/ZHBjHQENqW8/00053
+id04656/LDuq2UPHKoA/00157 id04656/LDuq2UPHKoA/00157
+id01509/UZL8Obdt--8/00181 id01509/UZL8Obdt--8/00181
+id05816/7jt8zGB27QQ/00017 id05816/7jt8zGB27QQ/00017
+id08456/7PKsuBS5LQI/00050 id08456/7PKsuBS5LQI/00050
+id06913/Tx0vAZhSPuE/00077 id06913/Tx0vAZhSPuE/00077
+id02465/UEmI4r5G-5Y/00117 id02465/UEmI4r5G-5Y/00117
+id01460/9sefvU9y4Kw/00046 id01460/9sefvU9y4Kw/00046
+id01567/uYDx0vIVy_A/00429 id01567/uYDx0vIVy_A/00429
+id07961/qott7SmhA-A/00351 id07961/qott7SmhA-A/00351
+id00866/Awi1Q0yib1s/00092 id00866/Awi1Q0yib1s/00092
+id02086/CqJKcn8m_Xo/00152 id02086/CqJKcn8m_Xo/00152
+id05015/Obbv73CqtmQ/00137 id05015/Obbv73CqtmQ/00137
+id01041/1UYZqPpavtk/00001 id01041/1UYZqPpavtk/00001
+id01593/GiLxqKSI68o/00188 id01593/GiLxqKSI68o/00188
+id02317/IR0psXbOjdc/00176 id02317/IR0psXbOjdc/00176
+id01066/X33aJxc3Kt0/00112 id01066/X33aJxc3Kt0/00112
+id08456/VU3fkD-QqPw/00206 id08456/VU3fkD-QqPw/00206
+id04536/wat5sbCSs0k/00470 id04536/wat5sbCSs0k/00470
+id01066/4KOSmyAMipc/00020 id01066/4KOSmyAMipc/00020
+id02445/f5u3ktNPHAk/00074 id02445/f5u3ktNPHAk/00074
+id03041/NJUcU7j30JI/00011 id03041/NJUcU7j30JI/00011
+id00817/vUezvJDh_tA/00394 id00817/vUezvJDh_tA/00394
+id04478/sw50KQMY8vw/00298 id04478/sw50KQMY8vw/00298
+id04657/hMrgeYf5ToQ/00267 id04657/hMrgeYf5ToQ/00267
+id02548/VdjlKRtLD_w/00206 id02548/VdjlKRtLD_w/00206
+id06310/4oJF1NW2bIg/00006 id06310/4oJF1NW2bIg/00006
+id01509/jqbtAt91alI/00329 id01509/jqbtAt91alI/00329
+id07414/oXx9CvIeFFY/00407 id07414/oXx9CvIeFFY/00407
+id04570/mwhiZtTZYX0/00271 id04570/mwhiZtTZYX0/00271
+id00812/AzDjo0Uyk4Y/00061 id00812/AzDjo0Uyk4Y/00061
+id05999/MJwLq17VoMA/00146 id05999/MJwLq17VoMA/00146
+id07414/dsqrI97WQHE/00319 id07414/dsqrI97WQHE/00319
+id05015/C3KsCD-pUgs/00046 id05015/C3KsCD-pUgs/00046
+id06484/Gh6H7Md_L2k/00053 id06484/Gh6H7Md_L2k/00053
+id00081/xlwJqdrzeMA/00291 id00081/xlwJqdrzeMA/00291
+id05055/RLN5nKfza4A/00219 id05055/RLN5nKfza4A/00219
+id05055/OKw_hph-hK8/00197 id05055/OKw_hph-hK8/00197
+id03839/xtBkY9xYpjA/00464 id03839/xtBkY9xYpjA/00464
+id07620/HEX00yF8LTs/00117 id07620/HEX00yF8LTs/00117
+id05816/hjrZgsKuvpw/00349 id05816/hjrZgsKuvpw/00349
+id02548/6LPbT49zy38/00050 id02548/6LPbT49zy38/00050
+id01000/7eYakM6qrTs/00006 id01000/7eYakM6qrTs/00006
+id02181/cNCj0pLxR24/00084 id02181/cNCj0pLxR24/00084
+id02086/sSliWvu6Ufs/00453 id02086/sSliWvu6Ufs/00453
+id03178/KHelFt1Jyyg/00057 id03178/KHelFt1Jyyg/00057
+id05594/8dYcSoUAQO8/00014 id05594/8dYcSoUAQO8/00014
+id05015/JmvJemqIeS0/00102 id05015/JmvJemqIeS0/00102
+id00081/EvCyt2keqW4/00065 id00081/EvCyt2keqW4/00065
+id07663/QWe7IIGrv5s/00146 id07663/QWe7IIGrv5s/00146
+id01618/kzxW2WAFWLI/00126 id01618/kzxW2WAFWLI/00126
+id00562/X7FJ3M3bz3c/00124 id00562/X7FJ3M3bz3c/00124
+id07961/bvPOvzukTE4/00224 id07961/bvPOvzukTE4/00224
+id03789/nv8sQplhvX0/00357 id03789/nv8sQplhvX0/00357
+id04295/VUHarbuO_eE/00125 id04295/VUHarbuO_eE/00125
+id01822/IaBziWYcwK4/00037 id01822/IaBziWYcwK4/00037
+id05015/X1opVctkTE8/00170 id05015/X1opVctkTE8/00170
+id01041/MMXznNig_iU/00248 id01041/MMXznNig_iU/00248
+id02465/EZ_F0hUZdS4/00054 id02465/EZ_F0hUZdS4/00054
+id04656/Bi7kCsbg5L0/00061 id04656/Bi7kCsbg5L0/00061
+id07494/K4ndWNAHgdU/00093 id07494/K4ndWNAHgdU/00093
+id07354/TKTT7fArInQ/00218 id07354/TKTT7fArInQ/00218
+id05714/Lu4PPvWXGn8/00014 id05714/Lu4PPvWXGn8/00014
+id05654/07pANazoyJg/00001 id05654/07pANazoyJg/00001
+id01066/FDp-ZLCWrIc/00054 id01066/FDp-ZLCWrIc/00054
+id05999/ZQJVmCJFjNs/00182 id05999/ZQJVmCJFjNs/00182
+id04570/5Fg6CLuRntk/00041 id04570/5Fg6CLuRntk/00041
+id08696/vqLNqYW4TQA/00476 id08696/vqLNqYW4TQA/00476
+id04862/2uYHadPvHRU/00016 id04862/2uYHadPvHRU/00016
+id03980/7MRUusImkno/00001 id03980/7MRUusImkno/00001
+id02542/QJKFnt1lHeE/00035 id02542/QJKFnt1lHeE/00035
+id04536/OYH-6uGB6jI/00322 id04536/OYH-6uGB6jI/00322
+id06484/dOTMnYZcY9Q/00126 id06484/dOTMnYZcY9Q/00126
+id04478/GZQGZOmFU5U/00063 id04478/GZQGZOmFU5U/00063
+id01224/tELp6C7FELU/00421 id01224/tELp6C7FELU/00421
+id03862/5m5iPZNJS6c/00022 id03862/5m5iPZNJS6c/00022
+id05124/lcDhSnyeN5E/00381 id05124/lcDhSnyeN5E/00381
+id08149/3V9V5sDAWTc/00001 id08149/3V9V5sDAWTc/00001
+id02181/iEF0MWApQms/00108 id02181/iEF0MWApQms/00108
+id04536/xrsxSF2qey8/00471 id04536/xrsxSF2qey8/00471
+id03178/9AJzTUwGbRk/00005 id03178/9AJzTUwGbRk/00005
+id01041/Izmh75CZNW0/00207 id01041/Izmh75CZNW0/00207
+id03041/g5YLpUZBNKc/00018 id03041/g5YLpUZBNKc/00018
+id03347/nSAKXYdEOOM/00297 id03347/nSAKXYdEOOM/00297
+id03347/pPWGEPixOoM/00337 id03347/pPWGEPixOoM/00337
+id07312/XBBpLMEjfUo/00048 id07312/XBBpLMEjfUo/00048
+id08456/6QFe7cYnZk4/00023 id08456/6QFe7cYnZk4/00023
+id05176/5Hk_hj0oXN8/00004 id05176/5Hk_hj0oXN8/00004
+id07426/DBBfi7aKLx4/00038 id07426/DBBfi7aKLx4/00038
+id07494/uhPKcTLLwcM/00347 id07494/uhPKcTLLwcM/00347
+id02576/agxjz_O2Wfs/00088 id02576/agxjz_O2Wfs/00088
+id01541/SvTz_Pn15Vk/00119 id01541/SvTz_Pn15Vk/00119
+id07414/Uxggn91FBog/00214 id07414/Uxggn91FBog/00214
+id04253/1HOlzefgLu8/00001 id04253/1HOlzefgLu8/00001
+id01567/RPUd0ua7RR0/00216 id01567/RPUd0ua7RR0/00216
+id04657/5DzZTPLgwTM/00044 id04657/5DzZTPLgwTM/00044
+id04006/zSMWS35kYdQ/00253 id04006/zSMWS35kYdQ/00253
+id03347/KT7B07WFWyM/00104 id03347/KT7B07WFWyM/00104
+id02445/z5u4yO1EsZo/00109 id02445/z5u4yO1EsZo/00109
+id00154/z1dLArSg5PQ/00190 id00154/z1dLArSg5PQ/00190
+id07414/Cn6Ws4oK1jg/00095 id07414/Cn6Ws4oK1jg/00095
+id02286/WHS1n7XUt_8/00103 id02286/WHS1n7XUt_8/00103
+id01509/Zmmnr4iRsCM/00230 id01509/Zmmnr4iRsCM/00230
+id04276/tGOA4fVnSgw/00448 id04276/tGOA4fVnSgw/00448
+id00419/nu9cRW2J4Dk/00420 id00419/nu9cRW2J4Dk/00420
+id07868/6RQX9l98N-g/00002 id07868/6RQX9l98N-g/00002
+id03839/1lh57VnuaKE/00004 id03839/1lh57VnuaKE/00004
+id03178/LT-BNQKA9NU/00075 id03178/LT-BNQKA9NU/00075
+id01460/Es6CkRmkIBY/00080 id01460/Es6CkRmkIBY/00080
+id06692/T2Xk7MO6m2g/00297 id06692/T2Xk7MO6m2g/00297
+id01892/d8b9y_CRE3M/00102 id01892/d8b9y_CRE3M/00102
+id07426/K_25cVSB-JU/00063 id07426/K_25cVSB-JU/00063
+id01333/LI6eLfuTn6I/00127 id01333/LI6eLfuTn6I/00127
+id00081/hIBFutPzn8s/00158 id00081/hIBFutPzn8s/00158
+id04536/2j8I_WX5mhY/00009 id04536/2j8I_WX5mhY/00009
+id04232/UElg0R7fmlk/00253 id04232/UElg0R7fmlk/00253
+id01460/eZR__GGkVw4/00221 id01460/eZR__GGkVw4/00221
+id01041/GymfYtTsKEU/00119 id01041/GymfYtTsKEU/00119
+id07396/xK1gClL60tY/00191 id07396/xK1gClL60tY/00191
+id05459/81o3ictaOnU/00075 id05459/81o3ictaOnU/00075
+id02685/yN8ilDTW-o4/00114 id02685/yN8ilDTW-o4/00114
+id02286/c8LjgwDQAkw/00137 id02286/c8LjgwDQAkw/00137
+id01541/SWcGs-DbV9Q/00100 id01541/SWcGs-DbV9Q/00100
+id01822/x4Fr2ceg_f8/00231 id01822/x4Fr2ceg_f8/00231
+id03347/FKY5V8wmX5k/00043 id03347/FKY5V8wmX5k/00043
+id00817/0GmSijZelGY/00001 id00817/0GmSijZelGY/00001
+id06209/ahL3F1x5sE4/00091 id06209/ahL3F1x5sE4/00091
+id06692/4k3Eo5s1Rwo/00057 id06692/4k3Eo5s1Rwo/00057
+id09017/sduESYpj2-I/00297 id09017/sduESYpj2-I/00297
+id07354/grg37qaxKjI/00329 id07354/grg37qaxKjI/00329
+id07802/X8I5FN64_Oc/00199 id07802/X8I5FN64_Oc/00199
+id07494/JV5S_SUcHmI/00088 id07494/JV5S_SUcHmI/00088
+id03524/eHrI5bD8hSs/00282 id03524/eHrI5bD8hSs/00282
+id01460/HNjuGz9ayBk/00109 id01460/HNjuGz9ayBk/00109
+id04570/961AefP1-is/00056 id04570/961AefP1-is/00056
+id00419/749eTxP4Us8/00061 id00419/749eTxP4Us8/00061
+id00017/OLguY5ofUrY/00039 id00017/OLguY5ofUrY/00039
+id08392/RogKVSjaAH0/00293 id08392/RogKVSjaAH0/00293
+id01066/lI1wGa1UhEM/00205 id01066/lI1wGa1UhEM/00205
+id07621/zSdriAuJUKo/00485 id07621/zSdriAuJUKo/00485
+id03862/JBkaiUNeMmk/00166 id03862/JBkaiUNeMmk/00166
+id00017/E6aqL_Nc410/00027 id00017/E6aqL_Nc410/00027
+id03839/fi-g--cBwnU/00348 id03839/fi-g--cBwnU/00348
+id05654/eLztZmvnk-k/00095 id05654/eLztZmvnk-k/00095
+id02548/wF5HfFXZCBI/00349 id02548/wF5HfFXZCBI/00349
+id02576/LAipS5WJ29s/00075 id02576/LAipS5WJ29s/00075
+id06692/SEPs17_AkTI/00295 id06692/SEPs17_AkTI/00295
+id05459/kkaYxtBZnNo/00348 id05459/kkaYxtBZnNo/00348
+id04232/MEGVEqgGCME/00167 id04232/MEGVEqgGCME/00167
+id01989/8CUktsB_2bA/00031 id01989/8CUktsB_2bA/00031
+id01066/kqP_NZ1FRlM/00176 id01066/kqP_NZ1FRlM/00176
+id03382/ockh8KdXJP8/00059 id03382/ockh8KdXJP8/00059
+id01593/pO180haP_vo/00410 id01593/pO180haP_vo/00410
+id07396/nTQDZrnGXXY/00179 id07396/nTQDZrnGXXY/00179
+id03030/rg-VUeksKaU/00257 id03030/rg-VUeksKaU/00257
+id08911/IddDkZwRflE/00053 id08911/IddDkZwRflE/00053
+id02317/K2GT02zavxo/00193 id02317/K2GT02zavxo/00193
+id01298/5P4ldDRuo5c/00065 id01298/5P4ldDRuo5c/00065
+id01989/Evbf6fMJNmk/00060 id01989/Evbf6fMJNmk/00060
+id05124/fNJI2A0v8yI/00357 id05124/fNJI2A0v8yI/00357
+id02465/RLi2ItGherA/00098 id02465/RLi2ItGherA/00098
+id07868/qMNfMcG6sh0/00346 id07868/qMNfMcG6sh0/00346
+id04366/tmoYV4kPOGU/00246 id04366/tmoYV4kPOGU/00246
+id06484/_ZkoebnFkVA/00110 id06484/_ZkoebnFkVA/00110
+id04276/I9gCyrZWFn0/00097 id04276/I9gCyrZWFn0/00097
+id03978/IMn6f0iDOtE/00032 id03978/IMn6f0iDOtE/00032
+id00419/w_0sK8WuSsg/00472 id00419/w_0sK8WuSsg/00472
+id04478/RwcHXQ3MvsQ/00109 id04478/RwcHXQ3MvsQ/00109
+id08696/cUmyIjpOYlY/00360 id08696/cUmyIjpOYlY/00360
+id04366/DqBQx6AZ1Nk/00083 id04366/DqBQx6AZ1Nk/00083
+id05459/RhOon49C3g8/00201 id05459/RhOon49C3g8/00201
+id04656/OzgjshkHUiA/00166 id04656/OzgjshkHUiA/00166
+id03969/x38Sqv819yE/00110 id03969/x38Sqv819yE/00110
+id00061/0G9G9oyFHI8/00001 id00061/0G9G9oyFHI8/00001
+id06913/IreNhnVfTkQ/00043 id06913/IreNhnVfTkQ/00043
+id01618/NqYUgbuImpk/00096 id01618/NqYUgbuImpk/00096
+id08552/y05_B9NXizo/00237 id08552/y05_B9NXizo/00237
+id01460/zcTt06bjKuA/00365 id01460/zcTt06bjKuA/00365
+id00866/nI-zVYcQX40/00220 id00866/nI-zVYcQX40/00220
+id08374/9eMfNJiKBPQ/00056 id08374/9eMfNJiKBPQ/00056
+id03524/nKxz0LxKZ58/00344 id03524/nKxz0LxKZ58/00344
+id09017/A3CAugN2cjk/00021 id09017/A3CAugN2cjk/00021
+id02685/NtHmnSLaGCA/00036 id02685/NtHmnSLaGCA/00036
+id01224/atjwjz0vAk8/00213 id01224/atjwjz0vAk8/00213
+id07961/gvLf2DggTu0/00271 id07961/gvLf2DggTu0/00271
+id01567/CCs8rZLCdVw/00043 id01567/CCs8rZLCdVw/00043
+id03347/nbmPriSE9NY/00316 id03347/nbmPriSE9NY/00316
+id06104/snzG1OymFgs/00273 id06104/snzG1OymFgs/00273
+id02019/xsXm-MSuD-E/00290 id02019/xsXm-MSuD-E/00290
+id00061/VugwXDj1ka4/00088 id00061/VugwXDj1ka4/00088
+id01224/4z68GFZuYKU/00028 id01224/4z68GFZuYKU/00028
+id03839/ajkGXKUvTWY/00296 id03839/ajkGXKUvTWY/00296
+id07874/N7fMpS_yaF4/00047 id07874/N7fMpS_yaF4/00047
+id05124/fRhAX7v_R6A/00365 id05124/fRhAX7v_R6A/00365
+id02181/ci_22Oqhwtc/00088 id02181/ci_22Oqhwtc/00088
+id07414/njxmqS9ncTA/00399 id07414/njxmqS9ncTA/00399
+id05176/yEMRxKA0vSw/00101 id05176/yEMRxKA0vSw/00101
+id03862/VVaxYHNmtA8/00269 id03862/VVaxYHNmtA8/00269
+id07396/X6KkvYh6rPA/00148 id07396/X6KkvYh6rPA/00148
+id06310/TkxTnoic67U/00130 id06310/TkxTnoic67U/00130
+id08374/Yh9O9ETuF_0/00250 id08374/Yh9O9ETuF_0/00250
+id02317/5moKZXlJTEs/00058 id02317/5moKZXlJTEs/00058
+id04536/EDCwhtRFARA/00172 id04536/EDCwhtRFARA/00172
+id03789/pz1jGMsPY9M/00381 id03789/pz1jGMsPY9M/00381
+id03127/wzS06bKAZ48/00354 id03127/wzS06bKAZ48/00354
+id08911/wedpC4fN4YY/00096 id08911/wedpC4fN4YY/00096
+id01106/6SFpvp42pMA/00014 id01106/6SFpvp42pMA/00014
+id02465/6jp5YsZYtHI/00021 id02465/6jp5YsZYtHI/00021
+id01618/Ay_BKx5-JOc/00046 id01618/Ay_BKx5-JOc/00046
+id04478/x07vvSVm2Yo/00363 id04478/x07vvSVm2Yo/00363
+id01593/u5AgUWl3fFU/00437 id01593/u5AgUWl3fFU/00437
+id03030/IpwcoJajjJI/00124 id03030/IpwcoJajjJI/00124
+id01593/t9TUbyp3xfs/00423 id01593/t9TUbyp3xfs/00423
+id07414/hUxcsEMKssA/00320 id07414/hUxcsEMKssA/00320
+id04366/L-56A5RNeWg/00124 id04366/L-56A5RNeWg/00124
+id07961/3EPjXGhfst4/00001 id07961/3EPjXGhfst4/00001
+id00061/mMOd25Ag7XY/00239 id00061/mMOd25Ag7XY/00239
+id01567/RQMG0K5AchU/00218 id01567/RQMG0K5AchU/00218
+id08552/PL5vk3XeKRM/00114 id08552/PL5vk3XeKRM/00114
+id04862/eX3wAZ0yr7w/00260 id04862/eX3wAZ0yr7w/00260
+id02086/CBNOvx4Phxw/00146 id02086/CBNOvx4Phxw/00146
+id01228/3wAkCYQR3fQ/00011 id01228/3wAkCYQR3fQ/00011
+id06484/MXwPpo1Dg7U/00073 id06484/MXwPpo1Dg7U/00073
+id01460/9fJy9zGdESI/00045 id01460/9fJy9zGdESI/00045

dataset/filelists/voxceleb2_test_n_500_seed_797_cross.txt

@@ -0,0 +1,500 @@
+id05459/18XmQEiGLnQ/00001 id07961/3EPjXGhfst4/00001
+id03980/7MRUusImkno/00001 id08696/0H1PxInJCK0/00001
+id05654/07pANazoyJg/00001 id04570/0YMGn6BI9rg/00001
+id00817/0GmSijZelGY/00001 id07354/0NjekFZqaY0/00001
+id05202/2gnLcAbAoSc/00001 id00817/0GmSijZelGY/00001
+id03041/5CfnYwQCW48/00001 id07354/0NjekFZqaY0/00001
+id03980/7MRUusImkno/00001 id07621/0CiFdFegqZM/00001
+id05850/B8kp8ed48JE/00001 id04253/1HOlzefgLu8/00001
+id01298/2K5F6xG-Rbs/00001 id05816/1dyCBbJ94iw/00001
+id07494/0P1wPmgz0Bk/00001 id07621/0CiFdFegqZM/00001
+id06913/4Ug7aJemzpg/00001 id04030/7mXUMuo5_NE/00001
+id02286/4LAIxvdvguc/00001 id05850/B8kp8ed48JE/00001
+id02548/0pAkJZmlFqc/00001 id05459/18XmQEiGLnQ/00001
+id08456/29EhSZDqzas/00001 id04295/1fSjOItVYVg/00001
+id04295/1fSjOItVYVg/00001 id02685/4JDRxqYC0a4/00001
+id04276/5M8NmCwTHZ0/00001 id05654/07pANazoyJg/00001
+id03030/5wOxV1wAgqA/00001 id03041/5CfnYwQCW48/00001
+id04656/1tZYt8jey54/00001 id07961/3EPjXGhfst4/00001
+id03980/7MRUusImkno/00001 id04536/0f_Yi_1CoeM/00001
+id03980/7MRUusImkno/00001 id05202/2gnLcAbAoSc/00001
+id01298/2K5F6xG-Rbs/00001 id04862/0zJh2FMTaDE/00001
+id02548/0pAkJZmlFqc/00001 id04478/2grMtwdG93I/00001
+id02685/4JDRxqYC0a4/00001 id01892/3vKPgjwFjbo/00001
+id07494/0P1wPmgz0Bk/00001 id04656/1tZYt8jey54/00001
+id00812/1Xfgvdu7oDo/00001 id00926/2Nd7f1yNQzE/00001
+id07426/1KNFfOFEhyI/00001 id03030/5wOxV1wAgqA/00001
+id00866/03SSllwNkGk/00001 id00812/1Xfgvdu7oDo/00001
+id04570/0YMGn6BI9rg/00001 id01892/3vKPgjwFjbo/00001
+id03041/5CfnYwQCW48/00001 id04030/7mXUMuo5_NE/00001
+id07494/0P1wPmgz0Bk/00001 id00081/2xYrsnvtUWc/00001
+id08392/0fwuibKviJU/00001 id05015/0Cu3AvWWOFI/00001
+id06692/2ptBBNIZXtI/00001 id04536/0f_Yi_1CoeM/00001
+id04253/1HOlzefgLu8/00001 id06104/02L1L9RFAgI/00001
+id02725/37kUrf6RJdw/00001 id02685/4JDRxqYC0a4/00001
+id04006/113VkmVVz1Q/00001 id04119/1uH67UruKlE/00001
+id01567/1Lx_ZqrK1bM/00001 id04030/7mXUMuo5_NE/00001
+id02445/3Rnk8eja3TU/00001 id05816/1dyCBbJ94iw/00001
+id03041/5CfnYwQCW48/00001 id03347/4xXZ75_TeSM/00001
+id04570/0YMGn6BI9rg/00001 id02317/0q4X8kPTlEY/00001
+id07426/1KNFfOFEhyI/00001 id01822/0QcHowaLAF0/00001
+id02577/0euHS_r5JH4/00001 id02725/37kUrf6RJdw/00001
+id07354/0NjekFZqaY0/00001 id05459/18XmQEiGLnQ/00001
+id06692/2ptBBNIZXtI/00001 id05850/B8kp8ed48JE/00001
+id01822/0QcHowaLAF0/00001 id07961/3EPjXGhfst4/00001
+id04366/0iG2Ub9zETM/00001 id03347/4xXZ75_TeSM/00001
+id03030/5wOxV1wAgqA/00001 id03789/0kdVSujPa9g/00001
+id04366/0iG2Ub9zETM/00001 id02286/4LAIxvdvguc/00001
+id00926/2Nd7f1yNQzE/00001 id02548/0pAkJZmlFqc/00001
+id03030/5wOxV1wAgqA/00001 id01298/2K5F6xG-Rbs/00001
+id01892/3vKPgjwFjbo/00001 id02317/0q4X8kPTlEY/00001
+id05202/2gnLcAbAoSc/00001 id04253/1HOlzefgLu8/00001
+id05714/2gvpaZcvAY4/00001 id06692/2ptBBNIZXtI/00001
+id07621/0CiFdFegqZM/00001 id07802/0RUpqvi3sPU/00001
+id03030/5wOxV1wAgqA/00001 id02317/0q4X8kPTlEY/00001
+id01822/0QcHowaLAF0/00001 id02445/3Rnk8eja3TU/00001
+id07961/3EPjXGhfst4/00001 id05459/18XmQEiGLnQ/00001
+id04950/2n4sGPqU9M8/00001 id07426/1KNFfOFEhyI/00001
+id04862/0zJh2FMTaDE/00001 id02465/0Ocu8l1eAng/00001
+id06104/02L1L9RFAgI/00001 id07312/0LWllHGohPY/00001
+id07414/110UMQovTR0/00001 id06692/2ptBBNIZXtI/00001
+id05015/0Cu3AvWWOFI/00001 id08696/0H1PxInJCK0/00001
+id02181/02gIO4WrZLY/00001 id00812/1Xfgvdu7oDo/00001
+id08392/0fwuibKviJU/00001 id01041/1UYZqPpavtk/00001
+id03347/4xXZ75_TeSM/00001 id04950/2n4sGPqU9M8/00001
+id07312/0LWllHGohPY/00001 id04950/2n4sGPqU9M8/00001
+id05202/2gnLcAbAoSc/00001 id05654/07pANazoyJg/00001
+id01041/1UYZqPpavtk/00001 id02317/0q4X8kPTlEY/00001
+id02057/0xZU7Oi9nvM/00001 id03178/2CT-6fnBC_o/00001
+id04006/113VkmVVz1Q/00001 id00817/0GmSijZelGY/00001
+id05850/B8kp8ed48JE/00001 id01892/3vKPgjwFjbo/00001
+id08696/0H1PxInJCK0/00001 id06692/2ptBBNIZXtI/00001
+id02057/0xZU7Oi9nvM/00001 id01541/2P7hzPq5iDw/00001
+id04006/113VkmVVz1Q/00001 id02057/0xZU7Oi9nvM/00001
+id04276/5M8NmCwTHZ0/00001 id04570/0YMGn6BI9rg/00001
+id07868/5YYJq3fSbH8/00001 id03030/5wOxV1wAgqA/00001
+id00812/1Xfgvdu7oDo/00001 id00154/0hjW3eTGAy8/00001
+id06692/2ptBBNIZXtI/00001 id05594/0ohBiepcHWI/00001
+id04536/0f_Yi_1CoeM/00001 id05202/2gnLcAbAoSc/00001
+id06310/1IAgr_CRnuE/00001 id05816/1dyCBbJ94iw/00001
+id01541/2P7hzPq5iDw/00001 id00419/1zffAxBod_c/00001
+id07354/0NjekFZqaY0/00001 id00866/03SSllwNkGk/00001
+id03347/4xXZ75_TeSM/00001 id02577/0euHS_r5JH4/00001
+id04119/1uH67UruKlE/00001 id04006/113VkmVVz1Q/00001
+id05714/2gvpaZcvAY4/00001 id07961/3EPjXGhfst4/00001
+id06104/02L1L9RFAgI/00001 id03178/2CT-6fnBC_o/00001
+id07354/0NjekFZqaY0/00001 id02445/3Rnk8eja3TU/00001
+id04030/7mXUMuo5_NE/00001 id03030/5wOxV1wAgqA/00001
+id07312/0LWllHGohPY/00001 id04536/0f_Yi_1CoeM/00001
+id03839/1jWHvl2qCq0/00001 id07802/0RUpqvi3sPU/00001
+id07621/0CiFdFegqZM/00001 id05816/1dyCBbJ94iw/00001
+id03839/1jWHvl2qCq0/00001 id03980/7MRUusImkno/00001
+id03030/5wOxV1wAgqA/00001 id02445/3Rnk8eja3TU/00001
+id03862/0w8W8jp7MJk/00001 id04253/1HOlzefgLu8/00001
+id05714/2gvpaZcvAY4/00001 id04119/1uH67UruKlE/00001
+id08392/0fwuibKviJU/00001 id07868/5YYJq3fSbH8/00001
+id01298/2K5F6xG-Rbs/00001 id03030/5wOxV1wAgqA/00001
+id05459/18XmQEiGLnQ/00001 id00817/0GmSijZelGY/00001
+id05850/B8kp8ed48JE/00001 id06692/2ptBBNIZXtI/00001
+id04295/1fSjOItVYVg/00001 id08456/29EhSZDqzas/00001
+id04570/0YMGn6BI9rg/00001 id02057/0xZU7Oi9nvM/00001
+id01541/2P7hzPq5iDw/00001 id00817/0GmSijZelGY/00001
+id07426/1KNFfOFEhyI/00001 id07354/0NjekFZqaY0/00001
+id04253/1HOlzefgLu8/00001 id06209/2zM9EAPsZZQ/00001
+id05850/B8kp8ed48JE/00001 id08392/0fwuibKviJU/00001
+id07802/0RUpqvi3sPU/00001 id02465/0Ocu8l1eAng/00001
+id04119/1uH67UruKlE/00001 id04862/0zJh2FMTaDE/00001
+id01541/2P7hzPq5iDw/00001 id08696/0H1PxInJCK0/00001
+id08696/0H1PxInJCK0/00001 id07802/0RUpqvi3sPU/00001
+id01228/2TIFacjgehY/00001 id07621/0CiFdFegqZM/00001
+id03178/2CT-6fnBC_o/00001 id07868/5YYJq3fSbH8/00001
+id05654/07pANazoyJg/00001 id01298/2K5F6xG-Rbs/00001
+id01822/0QcHowaLAF0/00001 id02548/0pAkJZmlFqc/00001
+id01618/0iFlmfmWVlY/00001 id08696/0H1PxInJCK0/00001
+id00812/1Xfgvdu7oDo/00001 id08456/29EhSZDqzas/00001
+id05594/0ohBiepcHWI/00001 id07312/0LWllHGohPY/00001
+id05714/2gvpaZcvAY4/00001 id06104/02L1L9RFAgI/00001
+id02445/3Rnk8eja3TU/00001 id07426/1KNFfOFEhyI/00001
+id05714/2gvpaZcvAY4/00001 id00817/0GmSijZelGY/00001
+id08696/0H1PxInJCK0/00001 id02317/0q4X8kPTlEY/00001
+id04950/2n4sGPqU9M8/00001 id04478/2grMtwdG93I/00001
+id01228/2TIFacjgehY/00001 id07414/110UMQovTR0/00001
+id00926/2Nd7f1yNQzE/00001 id01541/2P7hzPq5iDw/00001
+id05714/2gvpaZcvAY4/00001 id06913/4Ug7aJemzpg/00001
+id01228/2TIFacjgehY/00001 id03862/0w8W8jp7MJk/00001
+id03030/5wOxV1wAgqA/00001 id05015/0Cu3AvWWOFI/00001
+id02548/0pAkJZmlFqc/00001 id06692/2ptBBNIZXtI/00001
+id05202/2gnLcAbAoSc/00001 id04119/1uH67UruKlE/00001
+id04656/1tZYt8jey54/00001 id07426/1KNFfOFEhyI/00001
+id07312/0LWllHGohPY/00001 id03980/7MRUusImkno/00001
+id04366/0iG2Ub9zETM/00001 id00817/0GmSijZelGY/00001
+id07961/3EPjXGhfst4/00001 id01228/2TIFacjgehY/00001
+id00154/0hjW3eTGAy8/00001 id04295/1fSjOItVYVg/00001
+id04478/2grMtwdG93I/00001 id00154/0hjW3eTGAy8/00001
+id04570/0YMGn6BI9rg/00001 id05202/2gnLcAbAoSc/00001
+id04478/2grMtwdG93I/00001 id06913/4Ug7aJemzpg/00001
+id06104/02L1L9RFAgI/00001 id04295/1fSjOItVYVg/00001
+id05816/1dyCBbJ94iw/00001 id08392/0fwuibKviJU/00001
+id00926/2Nd7f1yNQzE/00001 id04536/0f_Yi_1CoeM/00001
+id00926/2Nd7f1yNQzE/00001 id02181/02gIO4WrZLY/00001
+id05459/18XmQEiGLnQ/00001 id02317/0q4X8kPTlEY/00001
+id05594/0ohBiepcHWI/00001 id01228/2TIFacjgehY/00001
+id02181/02gIO4WrZLY/00001 id07312/0LWllHGohPY/00001
+id00154/0hjW3eTGAy8/00001 id03839/1jWHvl2qCq0/00001
+id04030/7mXUMuo5_NE/00001 id02725/37kUrf6RJdw/00001
+id04295/1fSjOItVYVg/00001 id05714/2gvpaZcvAY4/00001
+id02548/0pAkJZmlFqc/00001 id04570/0YMGn6BI9rg/00001
+id04478/2grMtwdG93I/00001 id00866/03SSllwNkGk/00001
+id03030/5wOxV1wAgqA/00001 id04366/0iG2Ub9zETM/00001
+id02685/4JDRxqYC0a4/00001 id07426/1KNFfOFEhyI/00001
+id07802/0RUpqvi3sPU/00001 id07312/0LWllHGohPY/00001
+id02317/0q4X8kPTlEY/00001 id01892/3vKPgjwFjbo/00001
+id00154/0hjW3eTGAy8/00001 id00866/03SSllwNkGk/00001
+id02181/02gIO4WrZLY/00001 id02685/4JDRxqYC0a4/00001
+id03178/2CT-6fnBC_o/00001 id05459/18XmQEiGLnQ/00001
+id00926/2Nd7f1yNQzE/00001 id05202/2gnLcAbAoSc/00001
+id03041/5CfnYwQCW48/00001 id03178/2CT-6fnBC_o/00001
+id05850/B8kp8ed48JE/00001 id04006/113VkmVVz1Q/00001
+id01822/0QcHowaLAF0/00001 id04570/0YMGn6BI9rg/00001
+id04478/2grMtwdG93I/00001 id03839/1jWHvl2qCq0/00001
+id01298/2K5F6xG-Rbs/00001 id01228/2TIFacjgehY/00001
+id06310/1IAgr_CRnuE/00001 id04006/113VkmVVz1Q/00001
+id00154/0hjW3eTGAy8/00001 id04006/113VkmVVz1Q/00001
+id05816/1dyCBbJ94iw/00001 id01041/1UYZqPpavtk/00001
+id04570/0YMGn6BI9rg/00001 id04862/0zJh2FMTaDE/00001
+id06913/4Ug7aJemzpg/00001 id04862/0zJh2FMTaDE/00001
+id03862/0w8W8jp7MJk/00001 id02465/0Ocu8l1eAng/00001
+id04253/1HOlzefgLu8/00001 id01567/1Lx_ZqrK1bM/00001
+id06209/2zM9EAPsZZQ/00001 id01298/2K5F6xG-Rbs/00001
+id01822/0QcHowaLAF0/00001 id01541/2P7hzPq5iDw/00001
+id07312/0LWllHGohPY/00001 id02317/0q4X8kPTlEY/00001
+id06692/2ptBBNIZXtI/00001 id02445/3Rnk8eja3TU/00001
+id07414/110UMQovTR0/00001 id00154/0hjW3eTGAy8/00001
+id04478/2grMtwdG93I/00001 id03347/4xXZ75_TeSM/00001
+id04656/1tZYt8jey54/00001 id07802/0RUpqvi3sPU/00001
+id03839/1jWHvl2qCq0/00001 id06310/1IAgr_CRnuE/00001
+id02057/0xZU7Oi9nvM/00001 id01228/2TIFacjgehY/00001
+id00081/2xYrsnvtUWc/00001 id02057/0xZU7Oi9nvM/00001
+id03862/0w8W8jp7MJk/00001 id01892/3vKPgjwFjbo/00001
+id04570/0YMGn6BI9rg/00001 id06913/4Ug7aJemzpg/00001
+id08392/0fwuibKviJU/00001 id01567/1Lx_ZqrK1bM/00001
+id00081/2xYrsnvtUWc/00001 id07494/0P1wPmgz0Bk/00001
+id04536/0f_Yi_1CoeM/00001 id00081/2xYrsnvtUWc/00001
+id03839/1jWHvl2qCq0/00001 id05850/B8kp8ed48JE/00001
+id07621/0CiFdFegqZM/00001 id08456/29EhSZDqzas/00001
+id01822/0QcHowaLAF0/00001 id07868/5YYJq3fSbH8/00001
+id05202/2gnLcAbAoSc/00001 id03178/2CT-6fnBC_o/00001
+id06692/2ptBBNIZXtI/00001 id06913/4Ug7aJemzpg/00001
+id01041/1UYZqPpavtk/00001 id03030/5wOxV1wAgqA/00001
+id07426/1KNFfOFEhyI/00001 id08456/29EhSZDqzas/00001
+id04478/2grMtwdG93I/00001 id02548/0pAkJZmlFqc/00001
+id08392/0fwuibKviJU/00001 id01298/2K5F6xG-Rbs/00001
+id03041/5CfnYwQCW48/00001 id08696/0H1PxInJCK0/00001
+id04366/0iG2Ub9zETM/00001 id07426/1KNFfOFEhyI/00001
+id04950/2n4sGPqU9M8/00001 id07494/0P1wPmgz0Bk/00001
+id01822/0QcHowaLAF0/00001 id08392/0fwuibKviJU/00001
+id02577/0euHS_r5JH4/00001 id06692/2ptBBNIZXtI/00001
+id04570/0YMGn6BI9rg/00001 id00866/03SSllwNkGk/00001
+id05850/B8kp8ed48JE/00001 id08456/29EhSZDqzas/00001
+id01618/0iFlmfmWVlY/00001 id01041/1UYZqPpavtk/00001
+id07414/110UMQovTR0/00001 id04536/0f_Yi_1CoeM/00001
+id02057/0xZU7Oi9nvM/00001 id06913/4Ug7aJemzpg/00001
+id04536/0f_Yi_1CoeM/00001 id01041/1UYZqPpavtk/00001
+id04030/7mXUMuo5_NE/00001 id05850/B8kp8ed48JE/00001
+id04656/1tZYt8jey54/00001 id05459/18XmQEiGLnQ/00001
+id03789/0kdVSujPa9g/00001 id02057/0xZU7Oi9nvM/00001
+id01041/1UYZqPpavtk/00001 id05594/0ohBiepcHWI/00001
+id07494/0P1wPmgz0Bk/00001 id04006/113VkmVVz1Q/00001
+id00812/1Xfgvdu7oDo/00001 id04295/1fSjOItVYVg/00001
+id01541/2P7hzPq5iDw/00001 id02465/0Ocu8l1eAng/00001
+id04862/0zJh2FMTaDE/00001 id05594/0ohBiepcHWI/00001
+id05714/2gvpaZcvAY4/00001 id02286/4LAIxvdvguc/00001
+id06209/2zM9EAPsZZQ/00001 id05816/1dyCBbJ94iw/00001
+id05850/B8kp8ed48JE/00001 id00866/03SSllwNkGk/00001
+id07494/0P1wPmgz0Bk/00001 id07312/0LWllHGohPY/00001
+id04366/0iG2Ub9zETM/00001 id04570/0YMGn6BI9rg/00001
+id00866/03SSllwNkGk/00001 id03347/4xXZ75_TeSM/00001
+id02445/3Rnk8eja3TU/00001 id07802/0RUpqvi3sPU/00001
+id08696/0H1PxInJCK0/00001 id06209/2zM9EAPsZZQ/00001
+id02445/3Rnk8eja3TU/00001 id07621/0CiFdFegqZM/00001
+id08392/0fwuibKviJU/00001 id05850/B8kp8ed48JE/00001
+id00419/1zffAxBod_c/00001 id01228/2TIFacjgehY/00001
+id07354/0NjekFZqaY0/00001 id01041/1UYZqPpavtk/00001
+id04570/0YMGn6BI9rg/00001 id03347/4xXZ75_TeSM/00001
+id01892/3vKPgjwFjbo/00001 id02445/3Rnk8eja3TU/00001
+id00081/2xYrsnvtUWc/00001 id05459/18XmQEiGLnQ/00001
+id06104/02L1L9RFAgI/00001 id04570/0YMGn6BI9rg/00001
+id07961/3EPjXGhfst4/00001 id05654/07pANazoyJg/00001
+id00926/2Nd7f1yNQzE/00001 id03839/1jWHvl2qCq0/00001
+id02181/02gIO4WrZLY/00001 id08696/0H1PxInJCK0/00001
+id07426/1KNFfOFEhyI/00001 id05459/18XmQEiGLnQ/00001
+id03041/5CfnYwQCW48/00001 id06104/02L1L9RFAgI/00001
+id01298/2K5F6xG-Rbs/00001 id01541/2P7hzPq5iDw/00001
+id04570/0YMGn6BI9rg/00001 id01618/0iFlmfmWVlY/00001
+id02685/4JDRxqYC0a4/00001 id02548/0pAkJZmlFqc/00001
+id01822/0QcHowaLAF0/00001 id07426/1KNFfOFEhyI/00001
+id07868/5YYJq3fSbH8/00001 id07494/0P1wPmgz0Bk/00001
+id07802/0RUpqvi3sPU/00001 id03041/5CfnYwQCW48/00001
+id04656/1tZYt8jey54/00001 id01541/2P7hzPq5iDw/00001
+id03347/4xXZ75_TeSM/00001 id02445/3Rnk8eja3TU/00001
+id02548/0pAkJZmlFqc/00001 id01298/2K5F6xG-Rbs/00001
+id07354/0NjekFZqaY0/00001 id07426/1KNFfOFEhyI/00001
+id03862/0w8W8jp7MJk/00001 id01298/2K5F6xG-Rbs/00001
+id04536/0f_Yi_1CoeM/00001 id02465/0Ocu8l1eAng/00001
+id00081/2xYrsnvtUWc/00001 id04366/0iG2Ub9zETM/00001
+id04950/2n4sGPqU9M8/00001 id01822/0QcHowaLAF0/00001
+id06692/2ptBBNIZXtI/00001 id03030/5wOxV1wAgqA/00001
+id07312/0LWllHGohPY/00001 id04478/2grMtwdG93I/00001
+id03862/0w8W8jp7MJk/00001 id03030/5wOxV1wAgqA/00001
+id00081/2xYrsnvtUWc/00001 id08392/0fwuibKviJU/00001
+id02317/0q4X8kPTlEY/00001 id00154/0hjW3eTGAy8/00001
+id05594/0ohBiepcHWI/00001 id04536/0f_Yi_1CoeM/00001
+id07868/5YYJq3fSbH8/00001 id03839/1jWHvl2qCq0/00001
+id02577/0euHS_r5JH4/00001 id06913/4Ug7aJemzpg/00001
+id08456/29EhSZDqzas/00001 id01541/2P7hzPq5iDw/00001
+id01567/1Lx_ZqrK1bM/00001 id04119/1uH67UruKlE/00001
+id04253/1HOlzefgLu8/00001 id01228/2TIFacjgehY/00001
+id02445/3Rnk8eja3TU/00001 id02685/4JDRxqYC0a4/00001
+id05015/0Cu3AvWWOFI/00001 id02465/0Ocu8l1eAng/00001
+id07494/0P1wPmgz0Bk/00001 id05714/2gvpaZcvAY4/00001
+id02548/0pAkJZmlFqc/00001 id04006/113VkmVVz1Q/00001
+id00866/03SSllwNkGk/00001 id02317/0q4X8kPTlEY/00001
+id07354/0NjekFZqaY0/00001 id04253/1HOlzefgLu8/00001
+id00812/1Xfgvdu7oDo/00001 id03030/5wOxV1wAgqA/00001
+id02465/0Ocu8l1eAng/00001 id07354/0NjekFZqaY0/00001
+id04276/5M8NmCwTHZ0/00001 id03862/0w8W8jp7MJk/00001
+id01567/1Lx_ZqrK1bM/00001 id04253/1HOlzefgLu8/00001
+id01618/0iFlmfmWVlY/00001 id06913/4Ug7aJemzpg/00001
+id03862/0w8W8jp7MJk/00001 id08392/0fwuibKviJU/00001
+id07961/3EPjXGhfst4/00001 id00154/0hjW3eTGAy8/00001
+id02577/0euHS_r5JH4/00001 id01228/2TIFacjgehY/00001
+id05654/07pANazoyJg/00001 id03041/5CfnYwQCW48/00001
+id03980/7MRUusImkno/00001 id08392/0fwuibKviJU/00001
+id03178/2CT-6fnBC_o/00001 id04295/1fSjOItVYVg/00001
+id02317/0q4X8kPTlEY/00001 id03347/4xXZ75_TeSM/00001
+id02548/0pAkJZmlFqc/00001 id07426/1KNFfOFEhyI/00001
+id03839/1jWHvl2qCq0/00001 id05654/07pANazoyJg/00001
+id02548/0pAkJZmlFqc/00001 id07868/5YYJq3fSbH8/00001
+id04570/0YMGn6BI9rg/00001 id01041/1UYZqPpavtk/00001
+id07414/110UMQovTR0/00001 id00419/1zffAxBod_c/00001
+id00154/0hjW3eTGAy8/00001 id01618/0iFlmfmWVlY/00001
+id07494/0P1wPmgz0Bk/00001 id05654/07pANazoyJg/00001
+id01822/0QcHowaLAF0/00001 id06310/1IAgr_CRnuE/00001
+id05015/0Cu3AvWWOFI/00001 id05459/18XmQEiGLnQ/00001
+id05816/1dyCBbJ94iw/00001 id02317/0q4X8kPTlEY/00001
+id01541/2P7hzPq5iDw/00001 id05816/1dyCBbJ94iw/00001
+id06104/02L1L9RFAgI/00001 id01892/3vKPgjwFjbo/00001
+id04862/0zJh2FMTaDE/00001 id05850/B8kp8ed48JE/00001
+id05202/2gnLcAbAoSc/00001 id04366/0iG2Ub9zETM/00001
+id02286/4LAIxvdvguc/00001 id02725/37kUrf6RJdw/00001
+id04276/5M8NmCwTHZ0/00001 id01541/2P7hzPq5iDw/00001
+id02057/0xZU7Oi9nvM/00001 id03862/0w8W8jp7MJk/00001
+id06104/02L1L9RFAgI/00001 id00419/1zffAxBod_c/00001
+id04950/2n4sGPqU9M8/00001 id02181/02gIO4WrZLY/00001
+id04478/2grMtwdG93I/00001 id02685/4JDRxqYC0a4/00001
+id04006/113VkmVVz1Q/00001 id00081/2xYrsnvtUWc/00001
+id06692/2ptBBNIZXtI/00001 id03347/4xXZ75_TeSM/00001
+id03030/5wOxV1wAgqA/00001 id02465/0Ocu8l1eAng/00001
+id07312/0LWllHGohPY/00001 id03839/1jWHvl2qCq0/00001
+id04950/2n4sGPqU9M8/00001 id05654/07pANazoyJg/00001
+id02465/0Ocu8l1eAng/00001 id01618/0iFlmfmWVlY/00001
+id00419/1zffAxBod_c/00001 id02181/02gIO4WrZLY/00001
+id07426/1KNFfOFEhyI/00001 id05202/2gnLcAbAoSc/00001
+id07621/0CiFdFegqZM/00001 id08696/0H1PxInJCK0/00001
+id04006/113VkmVVz1Q/00001 id08392/0fwuibKviJU/00001
+id04478/2grMtwdG93I/00001 id02445/3Rnk8eja3TU/00001
+id03347/4xXZ75_TeSM/00001 id00154/0hjW3eTGAy8/00001
+id07312/0LWllHGohPY/00001 id02181/02gIO4WrZLY/00001
+id06310/1IAgr_CRnuE/00001 id02057/0xZU7Oi9nvM/00001
+id04366/0iG2Ub9zETM/00001 id05654/07pANazoyJg/00001
+id00419/1zffAxBod_c/00001 id04570/0YMGn6BI9rg/00001
+id04862/0zJh2FMTaDE/00001 id03862/0w8W8jp7MJk/00001
+id04366/0iG2Ub9zETM/00001 id00154/0hjW3eTGAy8/00001
+id00866/03SSllwNkGk/00001 id00081/2xYrsnvtUWc/00001
+id01618/0iFlmfmWVlY/00001 id02725/37kUrf6RJdw/00001
+id01892/3vKPgjwFjbo/00001 id07621/0CiFdFegqZM/00001
+id05015/0Cu3AvWWOFI/00001 id00926/2Nd7f1yNQzE/00001
+id06913/4Ug7aJemzpg/00001 id03839/1jWHvl2qCq0/00001
+id07312/0LWllHGohPY/00001 id07802/0RUpqvi3sPU/00001
+id06104/02L1L9RFAgI/00001 id02465/0Ocu8l1eAng/00001
+id04295/1fSjOItVYVg/00001 id01298/2K5F6xG-Rbs/00001
+id00866/03SSllwNkGk/00001 id05714/2gvpaZcvAY4/00001
+id06104/02L1L9RFAgI/00001 id01541/2P7hzPq5iDw/00001
+id02445/3Rnk8eja3TU/00001 id03789/0kdVSujPa9g/00001
+id00081/2xYrsnvtUWc/00001 id05816/1dyCBbJ94iw/00001
+id02548/0pAkJZmlFqc/00001 id03030/5wOxV1wAgqA/00001
+id04276/5M8NmCwTHZ0/00001 id01041/1UYZqPpavtk/00001
+id06913/4Ug7aJemzpg/00001 id07868/5YYJq3fSbH8/00001
+id04656/1tZYt8jey54/00001 id06692/2ptBBNIZXtI/00001
+id07494/0P1wPmgz0Bk/00001 id08696/0H1PxInJCK0/00001
+id04119/1uH67UruKlE/00001 id02317/0q4X8kPTlEY/00001
+id00419/1zffAxBod_c/00001 id04862/0zJh2FMTaDE/00001
+id03862/0w8W8jp7MJk/00001 id02445/3Rnk8eja3TU/00001
+id01892/3vKPgjwFjbo/00001 id04862/0zJh2FMTaDE/00001
+id04950/2n4sGPqU9M8/00001 id01618/0iFlmfmWVlY/00001
+id01228/2TIFacjgehY/00001 id01298/2K5F6xG-Rbs/00001
+id01041/1UYZqPpavtk/00001 id07961/3EPjXGhfst4/00001
+id07802/0RUpqvi3sPU/00001 id06913/4Ug7aJemzpg/00001
+id04276/5M8NmCwTHZ0/00001 id03030/5wOxV1wAgqA/00001
+id01567/1Lx_ZqrK1bM/00001 id05459/18XmQEiGLnQ/00001
+id02465/0Ocu8l1eAng/00001 id02725/37kUrf6RJdw/00001
+id05816/1dyCBbJ94iw/00001 id02181/02gIO4WrZLY/00001
+id06913/4Ug7aJemzpg/00001 id04950/2n4sGPqU9M8/00001
+id04276/5M8NmCwTHZ0/00001 id04253/1HOlzefgLu8/00001
+id07414/110UMQovTR0/00001 id06209/2zM9EAPsZZQ/00001
+id06310/1IAgr_CRnuE/00001 id03839/1jWHvl2qCq0/00001
+id03347/4xXZ75_TeSM/00001 id04006/113VkmVVz1Q/00001
+id01541/2P7hzPq5iDw/00001 id04253/1HOlzefgLu8/00001
+id08456/29EhSZDqzas/00001 id07494/0P1wPmgz0Bk/00001
+id07621/0CiFdFegqZM/00001 id05594/0ohBiepcHWI/00001
+id02685/4JDRxqYC0a4/00001 id04536/0f_Yi_1CoeM/00001
+id02317/0q4X8kPTlEY/00001 id08696/0H1PxInJCK0/00001
+id04253/1HOlzefgLu8/00001 id01041/1UYZqPpavtk/00001
+id01041/1UYZqPpavtk/00001 id03178/2CT-6fnBC_o/00001
+id05654/07pANazoyJg/00001 id01892/3vKPgjwFjbo/00001
+id04862/0zJh2FMTaDE/00001 id06310/1IAgr_CRnuE/00001
+id01541/2P7hzPq5iDw/00001 id04478/2grMtwdG93I/00001
+id02445/3Rnk8eja3TU/00001 id02057/0xZU7Oi9nvM/00001
+id08392/0fwuibKviJU/00001 id04570/0YMGn6BI9rg/00001
+id06692/2ptBBNIZXtI/00001 id02057/0xZU7Oi9nvM/00001
+id04950/2n4sGPqU9M8/00001 id04862/0zJh2FMTaDE/00001
+id03862/0w8W8jp7MJk/00001 id07621/0CiFdFegqZM/00001
+id07312/0LWllHGohPY/00001 id04656/1tZYt8jey54/00001
+id02577/0euHS_r5JH4/00001 id00866/03SSllwNkGk/00001
+id01228/2TIFacjgehY/00001 id02685/4JDRxqYC0a4/00001
+id00081/2xYrsnvtUWc/00001 id00419/1zffAxBod_c/00001
+id00154/0hjW3eTGAy8/00001 id04656/1tZYt8jey54/00001
+id03839/1jWHvl2qCq0/00001 id01618/0iFlmfmWVlY/00001
+id03862/0w8W8jp7MJk/00001 id02286/4LAIxvdvguc/00001
+id06310/1IAgr_CRnuE/00001 id08456/29EhSZDqzas/00001
+id02317/0q4X8kPTlEY/00001 id04276/5M8NmCwTHZ0/00001
+id06913/4Ug7aJemzpg/00001 id04366/0iG2Ub9zETM/00001
+id06310/1IAgr_CRnuE/00001 id00926/2Nd7f1yNQzE/00001
+id01228/2TIFacjgehY/00001 id02181/02gIO4WrZLY/00001
+id07414/110UMQovTR0/00001 id05594/0ohBiepcHWI/00001
+id03980/7MRUusImkno/00001 id03178/2CT-6fnBC_o/00001
+id03347/4xXZ75_TeSM/00001 id04478/2grMtwdG93I/00001
+id06692/2ptBBNIZXtI/00001 id05459/18XmQEiGLnQ/00001
+id00154/0hjW3eTGAy8/00001 id02725/37kUrf6RJdw/00001
+id01228/2TIFacjgehY/00001 id04006/113VkmVVz1Q/00001
+id00866/03SSllwNkGk/00001 id00926/2Nd7f1yNQzE/00001
+id05594/0ohBiepcHWI/00001 id04006/113VkmVVz1Q/00001
+id04656/1tZYt8jey54/00001 id01822/0QcHowaLAF0/00001
+id07354/0NjekFZqaY0/00001 id04536/0f_Yi_1CoeM/00001
+id07354/0NjekFZqaY0/00001 id04656/1tZYt8jey54/00001
+id04366/0iG2Ub9zETM/00001 id02057/0xZU7Oi9nvM/00001
+id03789/0kdVSujPa9g/00001 id01822/0QcHowaLAF0/00001
+id07621/0CiFdFegqZM/00001 id03347/4xXZ75_TeSM/00001
+id04030/7mXUMuo5_NE/00001 id04366/0iG2Ub9zETM/00001
+id00812/1Xfgvdu7oDo/00001 id07354/0NjekFZqaY0/00001
+id04536/0f_Yi_1CoeM/00001 id07494/0P1wPmgz0Bk/00001
+id04536/0f_Yi_1CoeM/00001 id05816/1dyCBbJ94iw/00001
+id03862/0w8W8jp7MJk/00001 id07868/5YYJq3fSbH8/00001
+id02685/4JDRxqYC0a4/00001 id05459/18XmQEiGLnQ/00001
+id06209/2zM9EAPsZZQ/00001 id07426/1KNFfOFEhyI/00001
+id07426/1KNFfOFEhyI/00001 id02317/0q4X8kPTlEY/00001
+id00926/2Nd7f1yNQzE/00001 id05594/0ohBiepcHWI/00001
+id00154/0hjW3eTGAy8/00001 id04950/2n4sGPqU9M8/00001
+id03041/5CfnYwQCW48/00001 id01892/3vKPgjwFjbo/00001
+id00419/1zffAxBod_c/00001 id00866/03SSllwNkGk/00001
+id02725/37kUrf6RJdw/00001 id05202/2gnLcAbAoSc/00001
+id04656/1tZYt8jey54/00001 id06913/4Ug7aJemzpg/00001
+id03862/0w8W8jp7MJk/00001 id04006/113VkmVVz1Q/00001
+id00419/1zffAxBod_c/00001 id04030/7mXUMuo5_NE/00001
+id06692/2ptBBNIZXtI/00001 id01541/2P7hzPq5iDw/00001
+id07354/0NjekFZqaY0/00001 id03041/5CfnYwQCW48/00001
+id03347/4xXZ75_TeSM/00001 id07802/0RUpqvi3sPU/00001
+id07354/0NjekFZqaY0/00001 id01298/2K5F6xG-Rbs/00001
+id02725/37kUrf6RJdw/00001 id03980/7MRUusImkno/00001
+id01618/0iFlmfmWVlY/00001 id02445/3Rnk8eja3TU/00001
+id05816/1dyCBbJ94iw/00001 id00081/2xYrsnvtUWc/00001
+id07354/0NjekFZqaY0/00001 id04478/2grMtwdG93I/00001
+id03980/7MRUusImkno/00001 id04295/1fSjOItVYVg/00001
+id02548/0pAkJZmlFqc/00001 id00081/2xYrsnvtUWc/00001
+id05459/18XmQEiGLnQ/00001 id03347/4xXZ75_TeSM/00001
+id04570/0YMGn6BI9rg/00001 id04006/113VkmVVz1Q/00001
+id06209/2zM9EAPsZZQ/00001 id01041/1UYZqPpavtk/00001
+id01228/2TIFacjgehY/00001 id02317/0q4X8kPTlEY/00001
+id07802/0RUpqvi3sPU/00001 id01541/2P7hzPq5iDw/00001
+id04862/0zJh2FMTaDE/00001 id01892/3vKPgjwFjbo/00001
+id04253/1HOlzefgLu8/00001 id07802/0RUpqvi3sPU/00001
+id06692/2ptBBNIZXtI/00001 id02286/4LAIxvdvguc/00001
+id01228/2TIFacjgehY/00001 id07961/3EPjXGhfst4/00001
+id05714/2gvpaZcvAY4/00001 id00812/1Xfgvdu7oDo/00001
+id03789/0kdVSujPa9g/00001 id03862/0w8W8jp7MJk/00001
+id04295/1fSjOItVYVg/00001 id07868/5YYJq3fSbH8/00001
+id04276/5M8NmCwTHZ0/00001 id02057/0xZU7Oi9nvM/00001
+id02286/4LAIxvdvguc/00001 id03862/0w8W8jp7MJk/00001
+id04478/2grMtwdG93I/00001 id05816/1dyCBbJ94iw/00001
+id08456/29EhSZDqzas/00001 id02725/37kUrf6RJdw/00001
+id02577/0euHS_r5JH4/00001 id07961/3EPjXGhfst4/00001
+id01618/0iFlmfmWVlY/00001 id00812/1Xfgvdu7oDo/00001
+id07312/0LWllHGohPY/00001 id03789/0kdVSujPa9g/00001
+id02685/4JDRxqYC0a4/00001 id03839/1jWHvl2qCq0/00001
+id04030/7mXUMuo5_NE/00001 id07802/0RUpqvi3sPU/00001
+id01567/1Lx_ZqrK1bM/00001 id04478/2grMtwdG93I/00001
+id02577/0euHS_r5JH4/00001 id02548/0pAkJZmlFqc/00001
+id04536/0f_Yi_1CoeM/00001 id03030/5wOxV1wAgqA/00001
+id03347/4xXZ75_TeSM/00001 id00081/2xYrsnvtUWc/00001
+id03980/7MRUusImkno/00001 id06209/2zM9EAPsZZQ/00001
+id01567/1Lx_ZqrK1bM/00001 id00154/0hjW3eTGAy8/00001
+id06104/02L1L9RFAgI/00001 id02057/0xZU7Oi9nvM/00001
+id04570/0YMGn6BI9rg/00001 id03980/7MRUusImkno/00001
+id08456/29EhSZDqzas/00001 id02286/4LAIxvdvguc/00001
+id07312/0LWllHGohPY/00001 id04366/0iG2Ub9zETM/00001
+id05654/07pANazoyJg/00001 id07426/1KNFfOFEhyI/00001
+id03839/1jWHvl2qCq0/00001 id03347/4xXZ75_TeSM/00001
+id04536/0f_Yi_1CoeM/00001 id04478/2grMtwdG93I/00001
+id05816/1dyCBbJ94iw/00001 id04862/0zJh2FMTaDE/00001
+id04950/2n4sGPqU9M8/00001 id00817/0GmSijZelGY/00001
+id07426/1KNFfOFEhyI/00001 id04862/0zJh2FMTaDE/00001
+id05459/18XmQEiGLnQ/00001 id00812/1Xfgvdu7oDo/00001
+id00154/0hjW3eTGAy8/00001 id03178/2CT-6fnBC_o/00001
+id04295/1fSjOItVYVg/00001 id07312/0LWllHGohPY/00001
+id05594/0ohBiepcHWI/00001 id04862/0zJh2FMTaDE/00001
+id03347/4xXZ75_TeSM/00001 id01541/2P7hzPq5iDw/00001
+id04536/0f_Yi_1CoeM/00001 id02445/3Rnk8eja3TU/00001
+id03862/0w8W8jp7MJk/00001 id04030/7mXUMuo5_NE/00001
+id00154/0hjW3eTGAy8/00001 id01541/2P7hzPq5iDw/00001
+id06913/4Ug7aJemzpg/00001 id03347/4xXZ75_TeSM/00001
+id08696/0H1PxInJCK0/00001 id04478/2grMtwdG93I/00001
+id04366/0iG2Ub9zETM/00001 id02445/3Rnk8eja3TU/00001
+id07354/0NjekFZqaY0/00001 id01567/1Lx_ZqrK1bM/00001
+id06913/4Ug7aJemzpg/00001 id05202/2gnLcAbAoSc/00001
+id04862/0zJh2FMTaDE/00001 id08696/0H1PxInJCK0/00001
+id03178/2CT-6fnBC_o/00001 id02685/4JDRxqYC0a4/00001
+id01822/0QcHowaLAF0/00001 id04950/2n4sGPqU9M8/00001
+id00081/2xYrsnvtUWc/00001 id06913/4Ug7aJemzpg/00001
+id07868/5YYJq3fSbH8/00001 id02465/0Ocu8l1eAng/00001
+id02181/02gIO4WrZLY/00001 id03862/0w8W8jp7MJk/00001
+id07868/5YYJq3fSbH8/00001 id05202/2gnLcAbAoSc/00001
+id02286/4LAIxvdvguc/00001 id03178/2CT-6fnBC_o/00001
+id01298/2K5F6xG-Rbs/00001 id01618/0iFlmfmWVlY/00001
+id03980/7MRUusImkno/00001 id04006/113VkmVVz1Q/00001
+id03862/0w8W8jp7MJk/00001 id08456/29EhSZDqzas/00001
+id01567/1Lx_ZqrK1bM/00001 id03041/5CfnYwQCW48/00001
+id02465/0Ocu8l1eAng/00001 id00419/1zffAxBod_c/00001
+id04570/0YMGn6BI9rg/00001 id04295/1fSjOItVYVg/00001
+id03862/0w8W8jp7MJk/00001 id04295/1fSjOItVYVg/00001
+id03789/0kdVSujPa9g/00001 id00866/03SSllwNkGk/00001
+id05654/07pANazoyJg/00001 id00926/2Nd7f1yNQzE/00001
+id05850/B8kp8ed48JE/00001 id02685/4JDRxqYC0a4/00001
+id03347/4xXZ75_TeSM/00001 id08392/0fwuibKviJU/00001
+id00926/2Nd7f1yNQzE/00001 id07312/0LWllHGohPY/00001
+id05850/B8kp8ed48JE/00001 id01041/1UYZqPpavtk/00001
+id03030/5wOxV1wAgqA/00001 id06913/4Ug7aJemzpg/00001
+id02057/0xZU7Oi9nvM/00001 id01041/1UYZqPpavtk/00001
+id03030/5wOxV1wAgqA/00001 id01041/1UYZqPpavtk/00001
+id01618/0iFlmfmWVlY/00001 id04366/0iG2Ub9zETM/00001
+id06310/1IAgr_CRnuE/00001 id04119/1uH67UruKlE/00001
+id05594/0ohBiepcHWI/00001 id02317/0q4X8kPTlEY/00001
+id01228/2TIFacjgehY/00001 id04119/1uH67UruKlE/00001
+id02286/4LAIxvdvguc/00001 id02445/3Rnk8eja3TU/00001
+id04030/7mXUMuo5_NE/00001 id00419/1zffAxBod_c/00001
+id01298/2K5F6xG-Rbs/00001 id02445/3Rnk8eja3TU/00001
+id07802/0RUpqvi3sPU/00001 id04862/0zJh2FMTaDE/00001
+id04006/113VkmVVz1Q/00001 id03347/4xXZ75_TeSM/00001
+id02317/0q4X8kPTlEY/00001 id05850/B8kp8ed48JE/00001
+id08456/29EhSZDqzas/00001 id04656/1tZYt8jey54/00001
+id04656/1tZYt8jey54/00001 id05816/1dyCBbJ94iw/00001
+id05202/2gnLcAbAoSc/00001 id06209/2zM9EAPsZZQ/00001

face_detection/README.md

@@ -0,0 +1 @@
+The code for Face Detection in this folder has been taken from the wonderful [face_alignment](https://github.com/1adrianb/face-alignment) repository. This has been modified to take batches of faces at a time. 

face_detection/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+
+__author__ = """Adrian Bulat"""
+__email__ = 'adrian.bulat@nottingham.ac.uk'
+__version__ = '1.0.1'
+
+from .api import FaceAlignment, LandmarksType, NetworkSize

face_detection/api.py

@@ -0,0 +1,98 @@
+from __future__ import print_function
+import os
+import torch
+from torch.utils.model_zoo import load_url
+from enum import Enum
+import numpy as np
+import cv2
+try:
+    import urllib.request as request_file
+except BaseException:
+    import urllib as request_file
+
+from .models import FAN, ResNetDepth
+from .utils import *
+
+
+class LandmarksType(Enum):
+    """Enum class defining the type of landmarks to detect.
+
+    ``_2D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face
+    ``_2halfD`` - this points represent the projection of the 3D points into 3D
+    ``_3D`` - detect the points ``(x,y,z)``` in a 3D space
+
+    """
+    _2D = 1
+    _2halfD = 2
+    _3D = 3
+
+
+class NetworkSize(Enum):
+    # TINY = 1
+    # SMALL = 2
+    # MEDIUM = 3
+    LARGE = 4
+
+    def __new__(cls, value):
+        member = object.__new__(cls)
+        member._value_ = value
+        return member
+
+    def __int__(self):
+        return self.value
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+
+class FaceAlignment:
+    def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,
+                 device='cuda', flip_input=False, face_detector='sfd', verbose=False):
+        self.device = device
+        self.flip_input = flip_input
+        self.landmarks_type = landmarks_type
+        self.verbose = verbose
+
+        network_size = int(network_size)
+
+        if 'cuda' in device:
+            torch.backends.cudnn.benchmark = True
+
+        # Get the face detector
+        face_detector_module = __import__('face_detection.detection.' + face_detector,
+                                          globals(), locals(), [face_detector], 0)
+        self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose)
+
+    def get_detections_for_batch(self, images):
+        images = images[..., ::-1]
+        detected_faces = self.face_detector.detect_from_batch(images.copy())
+        results = []
+
+        for i, d in enumerate(detected_faces):
+            # print("Inside facedection:", i, len(d))
+            if len(d) == 0:
+                results.append(None)
+                continue
+            d = d[0]
+            d = np.clip(d, 0, None)
+            
+            x1, y1, x2, y2 = map(int, d[:-1])
+            results.append((x1, y1, x2, y2))
+
+        return results
+    
+    def get_all_detections_for_batch(self, images):
+        #for multiface facedetection
+        images = images[..., ::-1]
+        detected_faces = self.face_detector.detect_from_batch(images.copy())
+        results = []
+
+        for i, d in enumerate(detected_faces):
+            # print("Inside facedection:", i, len(d))
+            if len(d) == 0:
+                results.append(None)
+                continue
+            d = [np.clip(dd, 0, None) for dd in d]
+            # d = [map(int, dd[:-1]) for dd in d]
+            d = [[int(ddd) for ddd in dd[:-1]] for dd in d]
+            results.append(d)
+
+        return results

face_detection/detection/__init__.py

@@ -0,0 +1 @@
+from .core import FaceDetector

face_detection/detection/core.py

@@ -0,0 +1,130 @@
+import logging
+import glob
+from tqdm import tqdm
+import numpy as np
+import torch
+import cv2
+
+
+class FaceDetector(object):
+    """An abstract class representing a face detector.
+
+    Any other face detection implementation must subclass it. All subclasses
+    must implement ``detect_from_image``, that return a list of detected
+    bounding boxes. Optionally, for speed considerations detect from path is
+    recommended.
+    """
+
+    def __init__(self, device, verbose):
+        self.device = device
+        self.verbose = verbose
+
+        if verbose:
+            if 'cpu' in device:
+                logger = logging.getLogger(__name__)
+                logger.warning("Detection running on CPU, this may be potentially slow.")
+
+        if 'cpu' not in device and 'cuda' not in device:
+            if verbose:
+                logger.error("Expected values for device are: {cpu, cuda} but got: %s", device)
+            raise ValueError
+
+    def detect_from_image(self, tensor_or_path):
+        """Detects faces in a given image.
+
+        This function detects the faces present in a provided BGR(usually)
+        image. The input can be either the image itself or the path to it.
+
+        Arguments:
+            tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path
+            to an image or the image itself.
+
+        Example::
+
+            >>> path_to_image = 'data/image_01.jpg'
+            ...   detected_faces = detect_from_image(path_to_image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+            >>> image = cv2.imread(path_to_image)
+            ...   detected_faces = detect_from_image(image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+
+        """
+        raise NotImplementedError
+
+    def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True):
+        """Detects faces from all the images present in a given directory.
+
+        Arguments:
+            path {string} -- a string containing a path that points to the folder containing the images
+
+        Keyword Arguments:
+            extensions {list} -- list of string containing the extensions to be
+            consider in the following format: ``.extension_name`` (default:
+            {['.jpg', '.png']}) recursive {bool} -- option wherever to scan the
+            folder recursively (default: {False}) show_progress_bar {bool} --
+            display a progressbar (default: {True})
+
+        Example:
+        >>> directory = 'data'
+        ...   detected_faces = detect_from_directory(directory)
+        {A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]}
+
+        """
+        if self.verbose:
+            logger = logging.getLogger(__name__)
+
+        if len(extensions) == 0:
+            if self.verbose:
+                logger.error("Expected at list one extension, but none was received.")
+            raise ValueError
+
+        if self.verbose:
+            logger.info("Constructing the list of images.")
+        additional_pattern = '/**/*' if recursive else '/*'
+        files = []
+        for extension in extensions:
+            files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
+
+        if self.verbose:
+            logger.info("Finished searching for images. %s images found", len(files))
+            logger.info("Preparing to run the detection.")
+
+        predictions = {}
+        for image_path in tqdm(files, disable=not show_progress_bar):
+            if self.verbose:
+                logger.info("Running the face detector on image: %s", image_path)
+            predictions[image_path] = self.detect_from_image(image_path)
+
+        if self.verbose:
+            logger.info("The detector was successfully run on all %s images", len(files))
+
+        return predictions
+
+    @property
+    def reference_scale(self):
+        raise NotImplementedError
+
+    @property
+    def reference_x_shift(self):
+        raise NotImplementedError
+
+    @property
+    def reference_y_shift(self):
+        raise NotImplementedError
+
+    @staticmethod
+    def tensor_or_path_to_ndarray(tensor_or_path, rgb=True):
+        """Convert path (represented as a string) or torch.tensor to a numpy.ndarray
+
+        Arguments:
+            tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself
+        """
+        if isinstance(tensor_or_path, str):
+            return cv2.imread(tensor_or_path) if not rgb else cv2.imread(tensor_or_path)[..., ::-1]
+        elif torch.is_tensor(tensor_or_path):
+            # Call cpu in case its coming from cuda
+            return tensor_or_path.cpu().numpy()[..., ::-1].copy() if not rgb else tensor_or_path.cpu().numpy()
+        elif isinstance(tensor_or_path, np.ndarray):
+            return tensor_or_path[..., ::-1].copy() if not rgb else tensor_or_path
+        else:
+            raise TypeError

face_detection/detection/sfd/__init__.py

@@ -0,0 +1 @@
+from .sfd_detector import SFDDetector as FaceDetector

face_detection/detection/sfd/bbox.py

@@ -0,0 +1,129 @@
+from __future__ import print_function
+import os
+import sys
+import cv2
+import random
+import datetime
+import time
+import math
+import argparse
+import numpy as np
+import torch
+
+try:
+    from iou import IOU
+except BaseException:
+    # IOU cython speedup 10x
+    def IOU(ax1, ay1, ax2, ay2, bx1, by1, bx2, by2):
+        sa = abs((ax2 - ax1) * (ay2 - ay1))
+        sb = abs((bx2 - bx1) * (by2 - by1))
+        x1, y1 = max(ax1, bx1), max(ay1, by1)
+        x2, y2 = min(ax2, bx2), min(ay2, by2)
+        w = x2 - x1
+        h = y2 - y1
+        if w < 0 or h < 0:
+            return 0.0
+        else:
+            return 1.0 * w * h / (sa + sb - w * h)
+
+
+def bboxlog(x1, y1, x2, y2, axc, ayc, aww, ahh):
+    xc, yc, ww, hh = (x2 + x1) / 2, (y2 + y1) / 2, x2 - x1, y2 - y1
+    dx, dy = (xc - axc) / aww, (yc - ayc) / ahh
+    dw, dh = math.log(ww / aww), math.log(hh / ahh)
+    return dx, dy, dw, dh
+
+
+def bboxloginv(dx, dy, dw, dh, axc, ayc, aww, ahh):
+    xc, yc = dx * aww + axc, dy * ahh + ayc
+    ww, hh = math.exp(dw) * aww, math.exp(dh) * ahh
+    x1, x2, y1, y2 = xc - ww / 2, xc + ww / 2, yc - hh / 2, yc + hh / 2
+    return x1, y1, x2, y2
+
+
+def nms(dets, thresh):
+    if 0 == len(dets):
+        return []
+    x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]])
+        xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]])
+
+        w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1)
+        ovr = w * h / (areas[i] + areas[order[1:]] - w * h)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def encode(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = torch.log(g_wh) / variances[1]
+    # return target for smooth_l1_loss
+    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = torch.cat((
+        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+def batch_decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = torch.cat((
+        priors[:, :, :2] + loc[:, :, :2] * variances[0] * priors[:, :, 2:],
+        priors[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1])), 2)
+    boxes[:, :, :2] -= boxes[:, :, 2:] / 2
+    boxes[:, :, 2:] += boxes[:, :, :2]
+    return boxes

face_detection/detection/sfd/detect.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn.functional as F
+
+import os
+import sys
+import cv2
+import random
+import datetime
+import math
+import argparse
+import numpy as np
+
+import scipy.io as sio
+import zipfile
+from .net_s3fd import s3fd
+from .bbox import *
+
+
+def detect(net, img, device):
+    img = img - np.array([104, 117, 123])
+    img = img.transpose(2, 0, 1)
+    img = img.reshape((1,) + img.shape)
+
+    if 'cuda' in device:
+        torch.backends.cudnn.benchmark = True
+
+    img = torch.from_numpy(img).float().to(device)
+    BB, CC, HH, WW = img.size()
+    with torch.no_grad():
+        olist = net(img)
+
+    bboxlist = []
+    for i in range(len(olist) // 2):
+        olist[i * 2] = F.softmax(olist[i * 2], dim=1)
+    olist = [oelem.data.cpu() for oelem in olist]
+    for i in range(len(olist) // 2):
+        ocls, oreg = olist[i * 2], olist[i * 2 + 1]
+        FB, FC, FH, FW = ocls.size()  # feature map size
+        stride = 2**(i + 2)    # 4,8,16,32,64,128
+        anchor = stride * 4
+        poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
+        for Iindex, hindex, windex in poss:
+            axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
+            score = ocls[0, 1, hindex, windex]
+            loc = oreg[0, :, hindex, windex].contiguous().view(1, 4)
+            priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]])
+            variances = [0.1, 0.2]
+            box = decode(loc, priors, variances)
+            x1, y1, x2, y2 = box[0] * 1.0
+            # cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1)
+            bboxlist.append([x1, y1, x2, y2, score])
+    bboxlist = np.array(bboxlist)
+    if 0 == len(bboxlist):
+        bboxlist = np.zeros((1, 5))
+
+    return bboxlist
+
+def batch_detect(net, imgs, device):
+    imgs = imgs - np.array([104, 117, 123])
+    imgs = imgs.transpose(0, 3, 1, 2)
+
+    if 'cuda' in device:
+        torch.backends.cudnn.benchmark = True
+
+    imgs = torch.from_numpy(imgs).float().to(device)
+    BB, CC, HH, WW = imgs.size()
+    with torch.no_grad():
+        olist = net(imgs)
+
+    bboxlist = []
+    for i in range(len(olist) // 2):
+        olist[i * 2] = F.softmax(olist[i * 2], dim=1)
+    olist = [oelem.data.cpu() for oelem in olist]
+    for i in range(len(olist) // 2):
+        ocls, oreg = olist[i * 2], olist[i * 2 + 1]
+        FB, FC, FH, FW = ocls.size()  # feature map size
+        stride = 2**(i + 2)    # 4,8,16,32,64,128
+        anchor = stride * 4
+        poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
+        for Iindex, hindex, windex in poss:
+            axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
+            score = ocls[:, 1, hindex, windex]
+            loc = oreg[:, :, hindex, windex].contiguous().view(BB, 1, 4)
+            priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]]).view(1, 1, 4)
+            variances = [0.1, 0.2]
+            box = batch_decode(loc, priors, variances)
+            box = box[:, 0] * 1.0
+            # cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1)
+            bboxlist.append(torch.cat([box, score.unsqueeze(1)], 1).cpu().numpy())
+    bboxlist = np.array(bboxlist)
+    if 0 == len(bboxlist):
+        bboxlist = np.zeros((1, BB, 5))
+
+    return bboxlist
+
+def flip_detect(net, img, device):
+    img = cv2.flip(img, 1)
+    b = detect(net, img, device)
+
+    bboxlist = np.zeros(b.shape)
+    bboxlist[:, 0] = img.shape[1] - b[:, 2]
+    bboxlist[:, 1] = b[:, 1]
+    bboxlist[:, 2] = img.shape[1] - b[:, 0]
+    bboxlist[:, 3] = b[:, 3]
+    bboxlist[:, 4] = b[:, 4]
+    return bboxlist
+
+
+def pts_to_bb(pts):
+    min_x, min_y = np.min(pts, axis=0)
+    max_x, max_y = np.max(pts, axis=0)
+    return np.array([min_x, min_y, max_x, max_y])

face_detection/detection/sfd/net_s3fd.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class L2Norm(nn.Module):
+    def __init__(self, n_channels, scale=1.0):
+        super(L2Norm, self).__init__()
+        self.n_channels = n_channels
+        self.scale = scale
+        self.eps = 1e-10
+        self.weight = nn.Parameter(torch.Tensor(self.n_channels))
+        self.weight.data *= 0.0
+        self.weight.data += self.scale
+
+    def forward(self, x):
+        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
+        x = x / norm * self.weight.view(1, -1, 1, 1)
+        return x
+
+
+class s3fd(nn.Module):
+    def __init__(self):
+        super(s3fd, self).__init__()
+        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
+        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+
+        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+
+        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+
+        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+
+        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+
+        self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3)
+        self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0)
+
+        self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
+        self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
+
+        self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
+        self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+
+        self.conv3_3_norm = L2Norm(256, scale=10)
+        self.conv4_3_norm = L2Norm(512, scale=8)
+        self.conv5_3_norm = L2Norm(512, scale=5)
+
+        self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+        self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+        self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+        self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+
+        self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1)
+        self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1)
+        self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+        self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1)
+        self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        h = F.relu(self.conv1_1(x))
+        h = F.relu(self.conv1_2(h))
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv2_1(h))
+        h = F.relu(self.conv2_2(h))
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv3_1(h))
+        h = F.relu(self.conv3_2(h))
+        h = F.relu(self.conv3_3(h))
+        f3_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv4_1(h))
+        h = F.relu(self.conv4_2(h))
+        h = F.relu(self.conv4_3(h))
+        f4_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv5_1(h))
+        h = F.relu(self.conv5_2(h))
+        h = F.relu(self.conv5_3(h))
+        f5_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.fc6(h))
+        h = F.relu(self.fc7(h))
+        ffc7 = h
+        h = F.relu(self.conv6_1(h))
+        h = F.relu(self.conv6_2(h))
+        f6_2 = h
+        h = F.relu(self.conv7_1(h))
+        h = F.relu(self.conv7_2(h))
+        f7_2 = h
+
+        f3_3 = self.conv3_3_norm(f3_3)
+        f4_3 = self.conv4_3_norm(f4_3)
+        f5_3 = self.conv5_3_norm(f5_3)
+
+        cls1 = self.conv3_3_norm_mbox_conf(f3_3)
+        reg1 = self.conv3_3_norm_mbox_loc(f3_3)
+        cls2 = self.conv4_3_norm_mbox_conf(f4_3)
+        reg2 = self.conv4_3_norm_mbox_loc(f4_3)
+        cls3 = self.conv5_3_norm_mbox_conf(f5_3)
+        reg3 = self.conv5_3_norm_mbox_loc(f5_3)
+        cls4 = self.fc7_mbox_conf(ffc7)
+        reg4 = self.fc7_mbox_loc(ffc7)
+        cls5 = self.conv6_2_mbox_conf(f6_2)
+        reg5 = self.conv6_2_mbox_loc(f6_2)
+        cls6 = self.conv7_2_mbox_conf(f7_2)
+        reg6 = self.conv7_2_mbox_loc(f7_2)
+
+        # max-out background label
+        chunk = torch.chunk(cls1, 4, 1)
+        bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2])
+        cls1 = torch.cat([bmax, chunk[3]], dim=1)
+
+        return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]

face_detection/detection/sfd/sfd_detector.py

@@ -0,0 +1,59 @@
+import os
+import cv2
+from torch.utils.model_zoo import load_url
+
+from ..core import FaceDetector
+
+from .net_s3fd import s3fd
+from .bbox import *
+from .detect import *
+
+models_urls = {
+    's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth',
+}
+
+
+class SFDDetector(FaceDetector):
+    def __init__(self, device, path_to_detector=os.path.join(os.path.dirname(os.path.abspath(__file__)), 's3fd.pth'), verbose=False):
+        super(SFDDetector, self).__init__(device, verbose)
+
+        # Initialise the face detector
+        if not os.path.isfile(path_to_detector):
+            model_weights = load_url(models_urls['s3fd'])
+        else:
+            model_weights = torch.load(path_to_detector)
+
+        self.face_detector = s3fd()
+        self.face_detector.load_state_dict(model_weights)
+        self.face_detector.to(device)
+        self.face_detector.eval()
+
+    def detect_from_image(self, tensor_or_path):
+        image = self.tensor_or_path_to_ndarray(tensor_or_path)
+
+        bboxlist = detect(self.face_detector, image, device=self.device)
+        keep = nms(bboxlist, 0.3)
+        bboxlist = bboxlist[keep, :]
+        bboxlist = [x for x in bboxlist if x[-1] > 0.5]
+
+        return bboxlist
+
+    def detect_from_batch(self, images):
+        bboxlists = batch_detect(self.face_detector, images, device=self.device)
+        keeps = [nms(bboxlists[:, i, :], 0.3) for i in range(bboxlists.shape[1])]
+        bboxlists = [bboxlists[keep, i, :] for i, keep in enumerate(keeps)]
+        bboxlists = [[x for x in bboxlist if x[-1] > 0.5] for bboxlist in bboxlists]
+
+        return bboxlists
+
+    @property
+    def reference_scale(self):
+        return 195
+
+    @property
+    def reference_x_shift(self):
+        return 0
+
+    @property
+    def reference_y_shift(self):
+        return 0

face_detection/models.py

@@ -0,0 +1,261 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3,
+                     stride=strd, padding=padding, bias=bias)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_planes, out_planes):
+        super(ConvBlock, self).__init__()
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.conv1 = conv3x3(in_planes, int(out_planes / 2))
+        self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
+        self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4))
+        self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
+        self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4))
+
+        if in_planes != out_planes:
+            self.downsample = nn.Sequential(
+                nn.BatchNorm2d(in_planes),
+                nn.ReLU(True),
+                nn.Conv2d(in_planes, out_planes,
+                          kernel_size=1, stride=1, bias=False),
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        residual = x
+
+        out1 = self.bn1(x)
+        out1 = F.relu(out1, True)
+        out1 = self.conv1(out1)
+
+        out2 = self.bn2(out1)
+        out2 = F.relu(out2, True)
+        out2 = self.conv2(out2)
+
+        out3 = self.bn3(out2)
+        out3 = F.relu(out3, True)
+        out3 = self.conv3(out3)
+
+        out3 = torch.cat((out1, out2, out3), 1)
+
+        if self.downsample is not None:
+            residual = self.downsample(residual)
+
+        out3 += residual
+
+        return out3
+
+
+class Bottleneck(nn.Module):
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class HourGlass(nn.Module):
+    def __init__(self, num_modules, depth, num_features):
+        super(HourGlass, self).__init__()
+        self.num_modules = num_modules
+        self.depth = depth
+        self.features = num_features
+
+        self._generate_network(self.depth)
+
+    def _generate_network(self, level):
+        self.add_module('b1_' + str(level), ConvBlock(self.features, self.features))
+
+        self.add_module('b2_' + str(level), ConvBlock(self.features, self.features))
+
+        if level > 1:
+            self._generate_network(level - 1)
+        else:
+            self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features))
+
+        self.add_module('b3_' + str(level), ConvBlock(self.features, self.features))
+
+    def _forward(self, level, inp):
+        # Upper branch
+        up1 = inp
+        up1 = self._modules['b1_' + str(level)](up1)
+
+        # Lower branch
+        low1 = F.avg_pool2d(inp, 2, stride=2)
+        low1 = self._modules['b2_' + str(level)](low1)
+
+        if level > 1:
+            low2 = self._forward(level - 1, low1)
+        else:
+            low2 = low1
+            low2 = self._modules['b2_plus_' + str(level)](low2)
+
+        low3 = low2
+        low3 = self._modules['b3_' + str(level)](low3)
+
+        up2 = F.interpolate(low3, scale_factor=2, mode='nearest')
+
+        return up1 + up2
+
+    def forward(self, x):
+        return self._forward(self.depth, x)
+
+
+class FAN(nn.Module):
+
+    def __init__(self, num_modules=1):
+        super(FAN, self).__init__()
+        self.num_modules = num_modules
+
+        # Base part
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.conv2 = ConvBlock(64, 128)
+        self.conv3 = ConvBlock(128, 128)
+        self.conv4 = ConvBlock(128, 256)
+
+        # Stacking part
+        for hg_module in range(self.num_modules):
+            self.add_module('m' + str(hg_module), HourGlass(1, 4, 256))
+            self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256))
+            self.add_module('conv_last' + str(hg_module),
+                            nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+            self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
+            self.add_module('l' + str(hg_module), nn.Conv2d(256,
+                                                            68, kernel_size=1, stride=1, padding=0))
+
+            if hg_module < self.num_modules - 1:
+                self.add_module(
+                    'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+                self.add_module('al' + str(hg_module), nn.Conv2d(68,
+                                                                 256, kernel_size=1, stride=1, padding=0))
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)), True)
+        x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        previous = x
+
+        outputs = []
+        for i in range(self.num_modules):
+            hg = self._modules['m' + str(i)](previous)
+
+            ll = hg
+            ll = self._modules['top_m_' + str(i)](ll)
+
+            ll = F.relu(self._modules['bn_end' + str(i)]
+                        (self._modules['conv_last' + str(i)](ll)), True)
+
+            # Predict heatmaps
+            tmp_out = self._modules['l' + str(i)](ll)
+            outputs.append(tmp_out)
+
+            if i < self.num_modules - 1:
+                ll = self._modules['bl' + str(i)](ll)
+                tmp_out_ = self._modules['al' + str(i)](tmp_out)
+                previous = previous + ll + tmp_out_
+
+        return outputs
+
+
+class ResNetDepth(nn.Module):
+
+    def __init__(self, block=Bottleneck, layers=[3, 8, 36, 3], num_classes=68):
+        self.inplanes = 64
+        super(ResNetDepth, self).__init__()
+        self.conv1 = nn.Conv2d(3 + 68, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x

face_detection/utils.py

@@ -0,0 +1,313 @@
+from __future__ import print_function
+import os
+import sys
+import time
+import torch
+import math
+import numpy as np
+import cv2
+
+
+def _gaussian(
+        size=3, sigma=0.25, amplitude=1, normalize=False, width=None,
+        height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,
+        mean_vert=0.5):
+    # handle some defaults
+    if width is None:
+        width = size
+    if height is None:
+        height = size
+    if sigma_horz is None:
+        sigma_horz = sigma
+    if sigma_vert is None:
+        sigma_vert = sigma
+    center_x = mean_horz * width + 0.5
+    center_y = mean_vert * height + 0.5
+    gauss = np.empty((height, width), dtype=np.float32)
+    # generate kernel
+    for i in range(height):
+        for j in range(width):
+            gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (
+                sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))
+    if normalize:
+        gauss = gauss / np.sum(gauss)
+    return gauss
+
+
+def draw_gaussian(image, point, sigma):
+    # Check if the gaussian is inside
+    ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)]
+    br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)]
+    if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1):
+        return image
+    size = 6 * sigma + 1
+    g = _gaussian(size)
+    g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
+    g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
+    img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+    img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+    assert (g_x[0] > 0 and g_y[1] > 0)
+    image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]
+          ] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]
+    image[image > 1] = 1
+    return image
+
+
+def transform(point, center, scale, resolution, invert=False):
+    """Generate and affine transformation matrix.
+
+    Given a set of points, a center, a scale and a targer resolution, the
+    function generates and affine transformation matrix. If invert is ``True``
+    it will produce the inverse transformation.
+
+    Arguments:
+        point {torch.tensor} -- the input 2D point
+        center {torch.tensor or numpy.array} -- the center around which to perform the transformations
+        scale {float} -- the scale of the face/object
+        resolution {float} -- the output resolution
+
+    Keyword Arguments:
+        invert {bool} -- define wherever the function should produce the direct or the
+        inverse transformation matrix (default: {False})
+    """
+    _pt = torch.ones(3)
+    _pt[0] = point[0]
+    _pt[1] = point[1]
+
+    h = 200.0 * scale
+    t = torch.eye(3)
+    t[0, 0] = resolution / h
+    t[1, 1] = resolution / h
+    t[0, 2] = resolution * (-center[0] / h + 0.5)
+    t[1, 2] = resolution * (-center[1] / h + 0.5)
+
+    if invert:
+        t = torch.inverse(t)
+
+    new_point = (torch.matmul(t, _pt))[0:2]
+
+    return new_point.int()
+
+
+def crop(image, center, scale, resolution=256.0):
+    """Center crops an image or set of heatmaps
+
+    Arguments:
+        image {numpy.array} -- an rgb image
+        center {numpy.array} -- the center of the object, usually the same as of the bounding box
+        scale {float} -- scale of the face
+
+    Keyword Arguments:
+        resolution {float} -- the size of the output cropped image (default: {256.0})
+
+    Returns:
+        [type] -- [description]
+    """  # Crop around the center point
+    """ Crops the image around the center. Input is expected to be an np.ndarray """
+    ul = transform([1, 1], center, scale, resolution, True)
+    br = transform([resolution, resolution], center, scale, resolution, True)
+    # pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0)
+    if image.ndim > 2:
+        newDim = np.array([br[1] - ul[1], br[0] - ul[0],
+                           image.shape[2]], dtype=np.int32)
+        newImg = np.zeros(newDim, dtype=np.uint8)
+    else:
+        newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
+        newImg = np.zeros(newDim, dtype=np.uint8)
+    ht = image.shape[0]
+    wd = image.shape[1]
+    newX = np.array(
+        [max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
+    newY = np.array(
+        [max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
+    oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
+    oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
+    newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]
+           ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
+    newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)),
+                        interpolation=cv2.INTER_LINEAR)
+    return newImg
+
+
+def get_preds_fromhm(hm, center=None, scale=None):
+    """Obtain (x,y) coordinates given a set of N heatmaps. If the center
+    and the scale is provided the function will return the points also in
+    the original coordinate frame.
+
+    Arguments:
+        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
+
+    Keyword Arguments:
+        center {torch.tensor} -- the center of the bounding box (default: {None})
+        scale {float} -- face scale (default: {None})
+    """
+    max, idx = torch.max(
+        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+    idx += 1
+    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+    preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
+    preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
+
+    for i in range(preds.size(0)):
+        for j in range(preds.size(1)):
+            hm_ = hm[i, j, :]
+            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                diff = torch.FloatTensor(
+                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],
+                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+                preds[i, j].add_(diff.sign_().mul_(.25))
+
+    preds.add_(-.5)
+
+    preds_orig = torch.zeros(preds.size())
+    if center is not None and scale is not None:
+        for i in range(hm.size(0)):
+            for j in range(hm.size(1)):
+                preds_orig[i, j] = transform(
+                    preds[i, j], center, scale, hm.size(2), True)
+
+    return preds, preds_orig
+
+def get_preds_fromhm_batch(hm, centers=None, scales=None):
+    """Obtain (x,y) coordinates given a set of N heatmaps. If the centers
+    and the scales is provided the function will return the points also in
+    the original coordinate frame.
+
+    Arguments:
+        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
+
+    Keyword Arguments:
+        centers {torch.tensor} -- the centers of the bounding box (default: {None})
+        scales {float} -- face scales (default: {None})
+    """
+    max, idx = torch.max(
+        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+    idx += 1
+    preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+    preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
+    preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
+
+    for i in range(preds.size(0)):
+        for j in range(preds.size(1)):
+            hm_ = hm[i, j, :]
+            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                diff = torch.FloatTensor(
+                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],
+                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+                preds[i, j].add_(diff.sign_().mul_(.25))
+
+    preds.add_(-.5)
+
+    preds_orig = torch.zeros(preds.size())
+    if centers is not None and scales is not None:
+        for i in range(hm.size(0)):
+            for j in range(hm.size(1)):
+                preds_orig[i, j] = transform(
+                    preds[i, j], centers[i], scales[i], hm.size(2), True)
+
+    return preds, preds_orig
+
+def shuffle_lr(parts, pairs=None):
+    """Shuffle the points left-right according to the axis of symmetry
+    of the object.
+
+    Arguments:
+        parts {torch.tensor} -- a 3D or 4D object containing the
+        heatmaps.
+
+    Keyword Arguments:
+        pairs {list of integers} -- [order of the flipped points] (default: {None})
+    """
+    if pairs is None:
+        pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
+                 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35,
+                 34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41,
+                 40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63,
+                 62, 61, 60, 67, 66, 65]
+    if parts.ndimension() == 3:
+        parts = parts[pairs, ...]
+    else:
+        parts = parts[:, pairs, ...]
+
+    return parts
+
+
+def flip(tensor, is_label=False):
+    """Flip an image or a set of heatmaps left-right
+
+    Arguments:
+        tensor {numpy.array or torch.tensor} -- [the input image or heatmaps]
+
+    Keyword Arguments:
+        is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False})
+    """
+    if not torch.is_tensor(tensor):
+        tensor = torch.from_numpy(tensor)
+
+    if is_label:
+        tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1)
+    else:
+        tensor = tensor.flip(tensor.ndimension() - 1)
+
+    return tensor
+
+# From pyzolib/paths.py (https://bitbucket.org/pyzo/pyzolib/src/tip/paths.py)
+
+
+def appdata_dir(appname=None, roaming=False):
+    """ appdata_dir(appname=None, roaming=False)
+
+    Get the path to the application directory, where applications are allowed
+    to write user specific files (e.g. configurations). For non-user specific
+    data, consider using common_appdata_dir().
+    If appname is given, a subdir is appended (and created if necessary).
+    If roaming is True, will prefer a roaming directory (Windows Vista/7).
+    """
+
+    # Define default user directory
+    userDir = os.getenv('FACEALIGNMENT_USERDIR', None)
+    if userDir is None:
+        userDir = os.path.expanduser('~')
+        if not os.path.isdir(userDir):  # pragma: no cover
+            userDir = '/var/tmp'  # issue #54
+
+    # Get system app data dir
+    path = None
+    if sys.platform.startswith('win'):
+        path1, path2 = os.getenv('LOCALAPPDATA'), os.getenv('APPDATA')
+        path = (path2 or path1) if roaming else (path1 or path2)
+    elif sys.platform.startswith('darwin'):
+        path = os.path.join(userDir, 'Library', 'Application Support')
+    # On Linux and as fallback
+    if not (path and os.path.isdir(path)):
+        path = userDir
+
+    # Maybe we should store things local to the executable (in case of a
+    # portable distro or a frozen application that wants to be portable)
+    prefix = sys.prefix
+    if getattr(sys, 'frozen', None):
+        prefix = os.path.abspath(os.path.dirname(sys.executable))
+    for reldir in ('settings', '../settings'):
+        localpath = os.path.abspath(os.path.join(prefix, reldir))
+        if os.path.isdir(localpath):  # pragma: no cover
+            try:
+                open(os.path.join(localpath, 'test.write'), 'wb').close()
+                os.remove(os.path.join(localpath, 'test.write'))
+            except IOError:
+                pass  # We cannot write in this directory
+            else:
+                path = localpath
+                break
+
+    # Get path specific for this app
+    if appname:
+        if path == userDir:
+            appname = '.' + appname.lstrip('.')  # Make it a hidden directory
+        path = os.path.join(path, appname)
+        if not os.path.isdir(path):  # pragma: no cover
+            os.mkdir(path)
+
+    # Done
+    return path

generate.py

@@ -0,0 +1,398 @@
+''' consistent initial noise for video generation'''
+import cv2
+import os
+from os.path import join, basename, dirname, splitext
+import shutil
+import argparse
+import numpy as np
+import random
+import torch, torchvision
+import subprocess
+from audio import audio
+import face_detection
+from tqdm import tqdm
+
+from guided_diffusion import dist_util, logger
+from guided_diffusion.resample import create_named_schedule_sampler
+from guided_diffusion.script_util import (
+    tfg_model_and_diffusion_defaults,
+    tfg_create_model_and_diffusion,
+    args_to_dict,
+    add_dict_to_argparser,
+)
+
+from guided_diffusion.tfg_data_util import (
+    tfg_process_batch,
+) 
+
+def get_frame_id(frame):
+    return int(basename(frame).split('.')[0])
+
+def crop_audio_window(spec, start_frame, args ):
+    if type(start_frame) == int:
+        start_frame_num = start_frame
+    else:
+        start_frame_num = get_frame_id(start_frame)
+    start_idx = int(args.mel_steps_per_sec * (start_frame_num / float(args.video_fps)))
+    end_idx = start_idx + args.syncnet_mel_step_size
+    return spec[start_idx : end_idx, :]
+
+def load_all_indiv_mels(path, args):
+    in_path = path
+    out_dir = join(args.sample_path, "temp", basename(in_path).replace(".mp4", ""))
+    os.makedirs(out_dir, exist_ok= True)
+    out_path = join(out_dir, "audio.wav")
+    command2 = 'ffmpeg -loglevel error -y -i {} -strict -2 {}'.format(in_path, out_path)
+    subprocess.call(command2, shell=True)
+    wav = audio.load_wav(out_path, args.sample_rate)
+    orig_mel = audio.melspectrogram(wav).T
+
+    all_indiv_mels = []
+    # i=0
+    i=1
+    while True:
+        m = crop_audio_window(orig_mel.copy(), max(i - args.syncnet_T//2,0), args)
+        if (m.shape[0] != args.syncnet_mel_step_size):
+            break
+        all_indiv_mels.append(m.T)
+        i+=1
+    
+    #clean up
+    shutil.rmtree(join(args.sample_path, "temp"))
+    
+    return all_indiv_mels, wav
+
+def load_video_frames(path, args):
+    in_path = path
+    out_dir = join(args.sample_path, "temp", basename(in_path).replace(".mp4", ""), "image")
+    os.makedirs(out_dir, exist_ok= True)
+
+
+    command = "ffmpeg -loglevel error -y -i {} -vf fps={} -q:v 2 -qmin 1 {}/%05d.jpg".format(in_path, args.video_fps, out_dir)
+    subprocess.call(command, shell=True)
+
+    video_frames=[]
+    for i, img_name in enumerate(sorted(os.listdir(out_dir))):
+        img_path=join(out_dir, img_name)
+        img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
+        video_frames.append(img)
+
+    #clean up
+    shutil.rmtree(join(args.sample_path, "temp"))
+
+
+    return video_frames
+
+
+def get_smoothened_boxes(boxes, T):
+    for i in range(len(boxes)):
+        if i + T > len(boxes):
+            window = boxes[len(boxes) - T:]
+        else:
+            window = boxes[i : i + T]
+        boxes[i] = np.mean(window, axis=0)
+    return boxes
+
+def my_voxceleb2_crop(img):
+    return img[:-int(img.shape[0]*2.36/8) , int(img.shape[1]*1.8/8): -int(img.shape[1]*1.8/8)]
+
+def my_voxceleb2_crop_bboxs(img):
+    return 0,img.shape[0]-int(img.shape[0]*2.36/8), int(img.shape[1]*1.8/8), img.shape[1]-int(img.shape[1]*1.8/8)
+
+def face_detect(images, detector, args, resize=False):
+    batch_size = args.face_det_batch_size
+    
+    while 1:
+        predictions = []
+        try:
+            for i in range(0, len(images), batch_size):
+                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+        except RuntimeError:
+            if batch_size == 1:
+                raise RuntimeError('Image too big to run face detection on GPU')
+            batch_size //= 2
+            args.face_det_batch_size = batch_size
+            print('Recovering from OOM error; New batch size: {}'.format(batch_size))
+            continue
+        break
+
+    results = []
+    if type(args.pads) == str :
+        args.pads = [int(x) for x in args.pads.split(",")]
+    pady1, pady2, padx1, padx2 = args.pads
+    for rect, image in zip(predictions, images):
+        if rect is None:
+            raise ValueError('Face not detected!')
+
+        y1 = max(0, rect[1] - pady1)
+        y2 = min(image.shape[0], rect[3] + pady2)
+        x1 = max(0, rect[0] - padx1)
+        x2 = min(image.shape[1], rect[2] + padx2)
+        
+        results.append([x1, y1, x2, y2])
+
+    boxes = get_smoothened_boxes(np.array(results), T=5)
+    
+    if resize:
+        if args.is_voxceleb2:
+            results = [[cv2.resize(my_voxceleb2_crop(image),(args.image_size, args.image_size)), my_voxceleb2_crop_bboxs(image), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+        else:
+            results = [[cv2.resize(image[y1: y2, x1:x2],(args.image_size, args.image_size)), (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+    else:
+        results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+    return results 
+
+def normalise(tensor):
+    """ [-1,1]->[0,1]"""
+    return ((tensor+1)*0.5).clamp(0,1)
+
+def normalise2(tensor):
+    """ [0,1]->[-1,1]"""
+    return (tensor*2-1).clamp(-1,1)
+
+
+def sample_batch(batch, model, diffusion, args):
+    B, F, C, H, W = batch[f'image'].shape
+    sample_shape = (B*F, C, H, W)
+
+
+    #generate fixed noise
+    init_noise = None
+    if args.sampling_seed:
+
+        state = torch.get_rng_state()
+        torch.manual_seed(args.sampling_seed)
+        torch.cuda.manual_seed_all(args.sampling_seed)
+        init_noise = torch.randn((1,C,H,W))
+        #repeat noise for all frames
+        init_noise = init_noise.repeat(B*F,1,1,1)
+        torch.set_rng_state(state)
+
+
+    img_batch, model_kwargs = tfg_process_batch(batch, args.face_hide_percentage, 
+                                                use_ref=args.use_ref, 
+                                                use_audio=args.use_audio, 
+                                                # sampling_use_gt_for_ref=args.sampling_use_gt_for_ref, 
+                                                noise=init_noise)
+    
+
+    img_batch = img_batch.to(dist_util.dev())
+    model_kwargs = {k: v.to(dist_util.dev()) for k,v in model_kwargs.items()}
+    init_noise = init_noise.to(dist_util.dev()) if init_noise is not None else None
+
+    sample_fn = (
+        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
+    )
+    sample = sample_fn(
+        model,
+        sample_shape,
+        clip_denoised=args.clip_denoised,
+        model_kwargs=model_kwargs,
+        noise = init_noise
+    )
+    return sample, img_batch, model_kwargs
+
+
+def generate(video_path, audio_path, model, diffusion, detector,  args, out_path=None, save_orig=True):
+    video_frames = load_video_frames(video_path, args)
+    try:
+        face_det_results = face_detect(video_frames.copy(), detector, args, resize=True)
+    except Exception as e:
+        print("Error:", e, video_path, audio_path)
+        import traceback
+        print(traceback.format_exc())
+    wrong_all_indiv_mels, wrong_audio_wavform = load_all_indiv_mels(audio_path, args)
+
+    min_frames = min(len(video_frames), len(wrong_all_indiv_mels))
+    video_frames = video_frames[:min_frames]
+    face_det_results = face_det_results[:min_frames]
+    face_bboxes = [face_det_results[i][1] for i in range(min_frames)]
+    face_frames = torch.FloatTensor(np.transpose(np.asarray([face_det_results[i][0] for i in range(min_frames)], dtype=np.float32)/255.,(0,3,1,2)))#[N, C, H, W]
+    wrong_all_indiv_mels = torch.FloatTensor(np.asarray(wrong_all_indiv_mels[:min_frames])).unsqueeze(1) #[N, 1, h, w]
+
+    if save_orig:
+        if out_path is None:
+            out_path_orig = os.path.join(args.sample_path, splitext(basename(video_path))[0]+"_"+ splitext(basename(audio_path))[0]+"_orig.mp4") 
+        else:
+            out_path_orig = out_path.replace(".mp4", "_orig.mp4")
+        torchvision.io.write_video(
+            out_path_orig,
+            video_array=torch.from_numpy(np.array(video_frames)), fps = args.video_fps, video_codec='libx264',
+            audio_array=torch.from_numpy(wrong_audio_wavform).unsqueeze(0), audio_fps=args.sample_rate, audio_codec='aac'
+        ) 
+
+    if args.sampling_ref_type=='gt':
+        ref_frames = face_frames.clone()
+    elif args.sampling_ref_type=='first_frame':
+        ref_frames = face_frames[0:1].repeat(len(face_frames),1,1,1)
+    elif args.sampling_ref_type=='random':
+        rand_idx = random.Random(args.sampling_seed).randint(0, len(face_frames)-1)
+        ref_frames = face_frames[rand_idx:rand_idx+1].repeat(len(face_frames),1,1,1)
+
+    if args.sampling_input_type=='first_frame':
+        face_frames = face_frames[0:1].repeat(len(face_frames),1,1,1)
+        video_frames = np.array(video_frames[0:1]*len(video_frames))
+        face_bboxes = np.array(face_bboxes[0:1]*len(face_bboxes))
+
+        
+    generated_video_frames = []
+    b_s = args.sampling_batch_size
+    for i in range(0,min_frames, b_s*args.nframes):
+        video_frames_batch = video_frames[i:i+b_s*args.nframes]
+        face_bboxes_batch = face_bboxes[i:i+b_s*args.nframes]
+
+        try:
+            img_batch = face_frames[i:i+b_s*args.nframes] #[BF, C, H, W]  
+            img_batch = img_batch.reshape(-1, args.nframes, img_batch.size(-3), img_batch.size(-2), img_batch.size(-1))
+            ref_batch = ref_frames[i:i+b_s*args.nframes]
+            ref_batch = ref_batch.reshape(-1, args.nframes, ref_batch.size(-3), ref_batch.size(-2), ref_batch.size(-1))
+            wrong_indiv_mel_batch = wrong_all_indiv_mels[i:i+b_s*args.nframes] #[BF, 1, h, w]
+            wrong_indiv_mel_batch = wrong_indiv_mel_batch.reshape(-1, args.nframes, wrong_indiv_mel_batch.size(-3),wrong_indiv_mel_batch.size(-2),wrong_indiv_mel_batch.size(-1))
+        except: # of the last batch, if B*F % nframes!=0, then the above reshape throws error
+            # but internally everything is going to get converted to BF
+            # ie. (B,F, C, H, W) -> (B*F, C, H, W)  but (B*F, 1, C, H, W) -> (B*F, C, H, W) 
+            img_batch = face_frames[i:i+b_s*args.nframes] #[BF, C, H, W]
+            img_batch = img_batch.reshape(-1, 1, img_batch.size(-3), img_batch.size(-2), img_batch.size(-1))
+            ref_batch = ref_frames[i:i+b_s*args.nframes]
+            ref_batch = ref_batch.reshape(-1, 1, ref_batch.size(-3), ref_batch.size(-2), ref_batch.size(-1))
+            wrong_indiv_mel_batch = wrong_all_indiv_mels[i:i+b_s*args.nframes] #[BF, 1, h, w]
+            wrong_indiv_mel_batch = wrong_indiv_mel_batch.reshape(-1, 1, wrong_indiv_mel_batch.size(-3),wrong_indiv_mel_batch.size(-2),wrong_indiv_mel_batch.size(-1))
+        
+
+        batch = {"image":img_batch, 
+                "ref_img":ref_batch, 
+                "indiv_mels":wrong_indiv_mel_batch}   
+
+        sample, img_batch, model_kwargs = sample_batch(batch, model, diffusion, args)   
+        mask = model_kwargs['mask']
+        recon_batch = sample * mask + (1. -mask)*img_batch #[BF, C, H, W]
+        recon_batch = (normalise(recon_batch)*255).cpu().numpy().transpose(0,2,3,1) #[-1,1] -> [0,255]
+
+        for g,v,b in zip(recon_batch, video_frames_batch, face_bboxes_batch):
+            y1, y2, x1, x2 = b
+            g = cv2.resize(g.astype(np.uint8), (x2 - x1, y2 - y1))            
+            v[y1:y2, x1:x2] = g
+            generated_video_frames.append(v)
+
+    
+
+    print(wrong_audio_wavform.shape, np.array(generated_video_frames).shape)
+    min_time = len(generated_video_frames)/args.video_fps # because video is already smaller because it got chopped accoding to the mel array length
+    wrong_audio_wavform = wrong_audio_wavform[:int(min_time*args.sample_rate)]
+    print(wrong_audio_wavform.shape, np.array(generated_video_frames).shape)
+    if out_path is None:
+        out_path = os.path.join(args.sample_path, splitext(basename(video_path))[0]+"_"+ splitext(basename(audio_path))[0]+".mp4") 
+    torchvision.io.write_video(
+        out_path,
+        video_array=torch.from_numpy(np.array(generated_video_frames)), fps = args.video_fps, video_codec='libx264',
+        audio_array=torch.from_numpy(wrong_audio_wavform).unsqueeze(0), audio_fps=args.sample_rate, audio_codec='aac'
+    ) 
+
+
+
+              
+
+def generate_from_filelist(test_video_dir, filelist, model, diffusion, detector,  args):
+    video_names = []
+    audio_names = []
+    with open(filelist, "r") as f:
+        lines = f.readlines()
+    for line in tqdm(lines):
+        try:
+            audio_name, video_name = line.strip().split()
+            audio_path = join(test_video_dir, audio_name+'.mp4')
+            video_path = join(test_video_dir, video_name+'.mp4')
+            out_path = join(args.sample_path,audio_name.replace('/','.')+"_"+video_name.replace('/','.')+".mp4")
+            generate(video_path, audio_path, model, diffusion, detector,  args, out_path=out_path ,save_orig=args.save_orig)
+        except Exception as e:
+            print("Error:", e, video_path, audio_path)
+            import traceback
+            print(traceback.format_exc())
+
+        
+
+def main():
+    args = create_argparser().parse_args()
+    dist_util.setup_dist()
+    logger.configure(dir=args.sample_path, format_strs=["stdout", "log"])
+
+    logger.log("creating model...")
+    model, diffusion = tfg_create_model_and_diffusion(
+            **args_to_dict(args, tfg_model_and_diffusion_defaults().keys())
+    )
+    model.load_state_dict(
+            dist_util.load_state_dict(args.model_path, map_location='cpu')
+    )
+    model.to(dist_util.dev())
+    if args.use_fp16:
+        model.convert_to_fp16()
+    model.eval()
+
+    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, device='cuda' if torch.cuda.is_available() else 'cpu')
+
+    if args.generate_from_filelist:
+        generate_from_filelist(args.test_video_dir, args.filelist, model, diffusion, detector,  args)
+    else:
+        generate(args.video_path, args.audio_path, model, diffusion, detector,  args, out_path=args.out_path, save_orig=args.save_orig)
+
+
+def create_argparser():
+    defaults = dict(
+        # generate from a single audio-video pair
+        generate_from_filelist = False,
+        video_path = "",
+        audio_path = "",
+        out_path = None, 
+        save_orig = True,
+
+        #generate from filelist : generate_from_filelist = True
+        test_video_dir = "test_videos",
+        filelist = "test_filelist.txt",
+
+
+        use_fp16 = True,
+        #tfg specific
+        face_hide_percentage=0.5,
+        use_ref=False,
+        use_audio=False,
+        audio_as_style=False,
+        audio_as_style_encoder_mlp=False,
+        
+        #data args
+        nframes=1,
+        nrefer=0,
+        image_size=128,
+        syncnet_T = 5,
+        syncnet_mel_step_size = 16,
+        audio_frames_per_video = 16, #for tfg model, we use sound corresponding to 5 frames centred at that frame
+        audio_dim=80,
+        is_voxceleb2=True,
+
+        video_fps=25,
+        sample_rate=16000, #audio sampling rate
+        mel_steps_per_sec=80.,
+
+        #sampling args
+        clip_denoised=True, # not used in training
+        sampling_batch_size=2,
+        use_ddim=False,
+        model_path="",
+        sample_path="d2l_gen",
+        sample_partition="",
+        sampling_seed=None,
+        sampling_use_gt_for_ref=False,
+        sampling_ref_type='gt', #one of ['gt', 'first_frame', 'random']
+        sampling_input_type='gt', #one of ['gt', 'first_frame']
+        
+        # face detection args
+        face_det_batch_size=64,
+        pads = "0,0,0,0"
+    )
+    defaults.update(tfg_model_and_diffusion_defaults())
+    parser = argparse.ArgumentParser()
+    add_dict_to_argparser(parser, defaults)
+    return parser
+
+
+if __name__=="__main__":
+    main()

generate_dist.py

@@ -0,0 +1,428 @@
+''' consistent initial noise for video generation'''
+import cv2
+import os
+from os.path import join, basename, dirname, splitext
+import shutil
+import argparse
+import numpy as np
+import random
+import torch, torchvision
+import subprocess
+from audio import audio
+import face_detection
+from tqdm import tqdm
+from torch.nn.parallel.distributed import DistributedDataParallel as DDP
+from guided_diffusion import dist_util, logger
+from guided_diffusion.resample import create_named_schedule_sampler
+from guided_diffusion.script_util import (
+    tfg_model_and_diffusion_defaults,
+    tfg_create_model_and_diffusion,
+    args_to_dict,
+    add_dict_to_argparser,
+)
+from time import time
+import torch.distributed as dist
+from guided_diffusion.tfg_data_util import (
+    tfg_process_batch,
+) 
+
+def get_frame_id(frame):
+    return int(basename(frame).split('.')[0])
+
+def crop_audio_window(spec, start_frame, args ):
+    if type(start_frame) == int:
+        start_frame_num = start_frame
+    else:
+        start_frame_num = get_frame_id(start_frame)
+    start_idx = int(args.mel_steps_per_sec * (start_frame_num / float(args.video_fps)))
+    end_idx = start_idx + args.syncnet_mel_step_size
+    return spec[start_idx : end_idx, :]
+
+def load_all_indiv_mels(path, args):
+    in_path = path
+    out_dir = join(args.sample_path, "temp",str(dist.get_rank()), basename(in_path).replace(".mp4", ""))
+    os.makedirs(out_dir, exist_ok= True)
+    out_path = join(out_dir, "audio.wav")
+    command2 = 'ffmpeg -loglevel error -y -i {} -strict -2 {}'.format(in_path, out_path)
+    subprocess.call(command2, shell=True)
+    wav = audio.load_wav(out_path, args.sample_rate)
+    orig_mel = audio.melspectrogram(wav).T
+
+    all_indiv_mels = []
+    # i=0
+    i=1
+    while True:
+        m = crop_audio_window(orig_mel.copy(), max(i - args.syncnet_T//2,0), args)
+        if (m.shape[0] != args.syncnet_mel_step_size):
+            break
+        all_indiv_mels.append(m.T)
+        i+=1
+    
+    #clean up
+    shutil.rmtree(join(args.sample_path, "temp", str(dist.get_rank())))
+    
+    return all_indiv_mels, wav
+
+def load_video_frames(path, args):
+    in_path = path
+    out_dir = join(args.sample_path, "temp", str(dist.get_rank()), basename(in_path).replace(".mp4", ""), "image")
+    os.makedirs(out_dir, exist_ok= True)
+
+
+    command = "ffmpeg -loglevel error -y -i {} -vf fps={} -q:v 2 -qmin 1 {}/%05d.jpg".format(in_path, args.video_fps, out_dir)
+    subprocess.call(command, shell=True)
+
+    video_frames=[]
+    for i, img_name in enumerate(sorted(os.listdir(out_dir))):
+        img_path=join(out_dir, img_name)
+        img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
+        video_frames.append(img)
+
+    #clean up
+    shutil.rmtree(join(args.sample_path, "temp", str(dist.get_rank())))
+
+
+    return video_frames
+
+
+def get_smoothened_boxes(boxes, T):
+    for i in range(len(boxes)):
+        if i + T > len(boxes):
+            window = boxes[len(boxes) - T:]
+        else:
+            window = boxes[i : i + T]
+        boxes[i] = np.mean(window, axis=0)
+    return boxes
+
+def my_voxceleb2_crop(img):
+    return img[:-int(img.shape[0]*2.36/8) , int(img.shape[1]*1.8/8): -int(img.shape[1]*1.8/8)]
+
+def my_voxceleb2_crop_bboxs(img):
+    return 0,img.shape[0]-int(img.shape[0]*2.36/8), int(img.shape[1]*1.8/8), img.shape[1]-int(img.shape[1]*1.8/8)
+
+def face_detect(images, detector, args, resize=False):
+    batch_size = args.face_det_batch_size
+    
+    while 1:
+        predictions = []
+        try:
+            for i in range(0, len(images), batch_size):
+                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+        except RuntimeError:
+            if batch_size == 1:
+                raise RuntimeError('Image too big to run face detection on GPU')
+            batch_size //= 2
+            args.face_det_batch_size = batch_size
+            print('Recovering from OOM error; New batch size: {}'.format(batch_size))
+            continue
+        break
+
+    results = []
+    if type(args.pads) == str :
+        args.pads = [int(x) for x in args.pads.split(",")]
+    pady1, pady2, padx1, padx2 = args.pads
+    for rect, image in zip(predictions, images):
+        if rect is None:
+            raise ValueError('Face not detected!')
+
+        y1 = max(0, rect[1] - pady1)
+        y2 = min(image.shape[0], rect[3] + pady2)
+        x1 = max(0, rect[0] - padx1)
+        x2 = min(image.shape[1], rect[2] + padx2)
+        
+        results.append([x1, y1, x2, y2])
+
+    boxes = get_smoothened_boxes(np.array(results), T=5)
+    
+    if resize:
+        if args.is_voxceleb2:
+            results = [[cv2.resize(my_voxceleb2_crop(image),(args.image_size, args.image_size)), my_voxceleb2_crop_bboxs(image), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+        else:
+            results = [[cv2.resize(image[y1: y2, x1:x2],(args.image_size, args.image_size)), (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+    else:
+        results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+    return results 
+
+def normalise(tensor):
+    """ [-1,1]->[0,1]"""
+    return ((tensor+1)*0.5).clamp(0,1)
+
+def normalise2(tensor):
+    """ [0,1]->[-1,1]"""
+    return (tensor*2-1).clamp(-1,1)
+
+
+def sample_batch(batch, model, diffusion, args):
+    B, F, C, H, W = batch[f'image'].shape
+    sample_shape = (B*F, C, H, W)
+
+
+    #generate fixed noise
+    init_noise = None
+    if args.sampling_seed:
+
+        state = torch.get_rng_state()
+        torch.manual_seed(args.sampling_seed)
+        torch.cuda.manual_seed_all(args.sampling_seed)
+        init_noise = torch.randn((1,C,H,W))
+        #repeat noise for all frames
+        init_noise = init_noise.repeat(B*F,1,1,1)
+        torch.set_rng_state(state)
+
+
+    img_batch, model_kwargs = tfg_process_batch(batch, args.face_hide_percentage, 
+                                                use_ref=args.use_ref, 
+                                                use_audio=args.use_audio, 
+                                                # sampling_use_gt_for_ref=args.sampling_use_gt_for_ref, 
+                                                noise=init_noise)
+    
+
+    img_batch = img_batch.to(dist_util.dev())
+    model_kwargs = {k: v.to(dist_util.dev()) for k,v in model_kwargs.items()}
+    init_noise = init_noise.to(dist_util.dev()) if init_noise is not None else None
+
+    sample_fn = (
+        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
+    )
+    sample = sample_fn(
+        model,
+        sample_shape,
+        clip_denoised=args.clip_denoised,
+        model_kwargs=model_kwargs,
+        noise = init_noise
+    )
+    return sample, img_batch, model_kwargs
+
+
+def generate(video_path, audio_path, model, diffusion, detector,  args, out_path=None, save_orig=True):
+    video_frames = load_video_frames(video_path, args)
+    try:
+        face_det_results = face_detect(video_frames.copy(), detector, args, resize=True)
+    except Exception as e:
+        print("Error:", e, video_path, audio_path)
+        import traceback
+        print(traceback.format_exc())
+    wrong_all_indiv_mels, wrong_audio_wavform = load_all_indiv_mels(audio_path, args)
+
+    min_frames = min(len(video_frames), len(wrong_all_indiv_mels))
+    video_frames = video_frames[:min_frames]
+    face_det_results = face_det_results[:min_frames]
+    face_bboxes = [face_det_results[i][1] for i in range(min_frames)]
+    face_frames = torch.FloatTensor(np.transpose(np.asarray([face_det_results[i][0] for i in range(min_frames)], dtype=np.float32)/255.,(0,3,1,2)))#[N, C, H, W]
+    wrong_all_indiv_mels = torch.FloatTensor(np.asarray(wrong_all_indiv_mels[:min_frames])).unsqueeze(1) #[N, 1, h, w]
+
+    if save_orig:
+        if out_path is None:
+            out_path_orig = os.path.join(args.sample_path, splitext(basename(video_path))[0]+"_"+ splitext(basename(audio_path))[0]+"_orig.mp4") 
+        else:
+            out_path_orig = out_path.replace(".mp4", "_orig.mp4")
+        torchvision.io.write_video(
+            out_path_orig,
+            video_array=torch.from_numpy(np.array(video_frames)), fps = args.video_fps, video_codec='libx264',
+            audio_array=torch.from_numpy(wrong_audio_wavform).unsqueeze(0), audio_fps=args.sample_rate, audio_codec='aac'
+        ) 
+
+    if args.sampling_ref_type=='gt':
+        ref_frames = face_frames.clone()
+    elif args.sampling_ref_type=='first_frame':
+        ref_frames = face_frames[0:1].repeat(len(face_frames),1,1,1)
+    elif args.sampling_ref_type=='random':
+        rand_idx = random.Random(args.sampling_seed).randint(0, len(face_frames)-1)
+        ref_frames = face_frames[rand_idx:rand_idx+1].repeat(len(face_frames),1,1,1)
+
+    if args.sampling_input_type=='first_frame':
+        face_frames = face_frames[0:1].repeat(len(face_frames),1,1,1)
+        video_frames = np.array(video_frames[0:1]*len(video_frames))
+        face_bboxes = np.array(face_bboxes[0:1]*len(face_bboxes))
+
+
+    rank = dist.get_rank() 
+    world_size = dist.get_world_size()
+    chunk_size = int(np.ceil(min_frames/world_size))
+    start_idx = rank * chunk_size
+    end_idx = min(start_idx + chunk_size, min_frames)
+    generated_video_frames = []
+    b_s = args.sampling_batch_size
+
+    # print(rank,"/",world_size, "chunk: [",start_idx,"-", end_idx,"/",min_frames,"]")
+
+    dist.barrier()
+    torch.cuda.synchronize()
+    t1=time()
+    # for i in range(0,min_frames, b_s*args.nframes):
+    for i in range(start_idx,end_idx, b_s*args.nframes):
+        slice_end = min(i+b_s*args.nframes, end_idx)
+        # if rank==0:
+        #     print("rank 0: slice:",i,":",slice_end)
+        video_frames_batch = video_frames[i:slice_end]
+        face_bboxes_batch = face_bboxes[i:slice_end]
+
+        # try:
+        if  (slice_end-i) % args.nframes==0:
+            img_batch = face_frames[i:slice_end] #[BF, C, H, W]  
+            img_batch = img_batch.reshape(-1, args.nframes, img_batch.size(-3), img_batch.size(-2), img_batch.size(-1))
+            ref_batch = ref_frames[i:slice_end]
+            ref_batch = ref_batch.reshape(-1, args.nframes, ref_batch.size(-3), ref_batch.size(-2), ref_batch.size(-1))
+            wrong_indiv_mel_batch = wrong_all_indiv_mels[i:slice_end] #[BF, 1, h, w]
+            wrong_indiv_mel_batch = wrong_indiv_mel_batch.reshape(-1, args.nframes, wrong_indiv_mel_batch.size(-3),wrong_indiv_mel_batch.size(-2),wrong_indiv_mel_batch.size(-1))
+        # except: 
+        else: # of the last batch, if B*F % nframes!=0, then the above reshape throws error
+            # but internally everything is going to get converted to BF
+            # ie. (B,F, C, H, W) -> (B*F, C, H, W)  but (B*F, 1, C, H, W) -> (B*F, C, H, W) 
+            img_batch = face_frames[i:slice_end] #[BF, C, H, W]
+            img_batch = img_batch.reshape(-1, 1, img_batch.size(-3), img_batch.size(-2), img_batch.size(-1))
+            ref_batch = ref_frames[i:slice_end]
+            ref_batch = ref_batch.reshape(-1, 1, ref_batch.size(-3), ref_batch.size(-2), ref_batch.size(-1))
+            wrong_indiv_mel_batch = wrong_all_indiv_mels[i:slice_end] #[BF, 1, h, w]
+            wrong_indiv_mel_batch = wrong_indiv_mel_batch.reshape(-1, 1, wrong_indiv_mel_batch.size(-3),wrong_indiv_mel_batch.size(-2),wrong_indiv_mel_batch.size(-1))
+        
+
+        batch = {"image":img_batch, 
+                "ref_img":ref_batch, 
+                "indiv_mels":wrong_indiv_mel_batch}   
+
+        sample, img_batch, model_kwargs = sample_batch(batch, model, diffusion, args)   
+        mask = model_kwargs['mask']
+        recon_batch = sample * mask + (1. -mask)*img_batch #[BF, C, H, W]
+        recon_batch = (normalise(recon_batch)*255).cpu().numpy().transpose(0,2,3,1) #[-1,1] -> [0,255]
+
+        for g,v,b in zip(recon_batch, video_frames_batch, face_bboxes_batch):
+            y1, y2, x1, x2 = b
+            g = cv2.resize(g.astype(np.uint8), (x2 - x1, y2 - y1))            
+            v[y1:y2, x1:x2] = g
+            generated_video_frames.append(v)
+
+    torch.cuda.synchronize()
+    t3=time()
+    all_generated_video_frames = [None for _ in range(dist.get_world_size())]
+    dist.all_gather_object(all_generated_video_frames, generated_video_frames)  # gather not supported with NCCL
+    all_generated_video_frames_combined = []
+    [all_generated_video_frames_combined.extend(gvf) for gvf in all_generated_video_frames]
+    generated_video_frames = all_generated_video_frames_combined
+    
+    torch.cuda.synchronize()
+    t2=time()
+
+    if dist.get_rank() == 0:
+        print("Time taken for sampling, ", t2-t1, ",time without all  gather, ", t3-t1, ",frames/gpu, ", len(generated_video_frames), ",total frames, ", min_frames)
+        print(wrong_audio_wavform.shape, np.array(generated_video_frames).shape)
+        min_time = len(generated_video_frames)/args.video_fps # because video is already smaller because it got chopped accoding to the mel array length
+        wrong_audio_wavform = wrong_audio_wavform[:int(min_time*args.sample_rate)]
+        print(wrong_audio_wavform.shape, np.array(generated_video_frames).shape)
+        if out_path is None:
+            out_path = os.path.join(args.sample_path, splitext(basename(video_path))[0]+"_"+ splitext(basename(audio_path))[0]+".mp4") 
+        torchvision.io.write_video(
+            out_path,
+            video_array=torch.from_numpy(np.array(generated_video_frames)), fps = args.video_fps, video_codec='libx264',
+            audio_array=torch.from_numpy(wrong_audio_wavform).unsqueeze(0), audio_fps=args.sample_rate, audio_codec='aac'
+        ) 
+    dist.barrier()
+
+
+
+              
+
+def generate_from_filelist(test_video_dir, filelist, model, diffusion, detector,  args):
+    video_names = []
+    audio_names = []
+    with open(filelist, "r") as f:
+        lines = f.readlines()
+    for line in tqdm(lines):
+        try:
+            audio_name, video_name = line.strip().split()
+            audio_path = join(test_video_dir, audio_name+'.mp4')
+            video_path = join(test_video_dir, video_name+'.mp4')
+            out_path = join(args.sample_path,audio_name.replace('/','.')+"_"+video_name.replace('/','.')+".mp4")
+            generate(video_path, audio_path, model, diffusion, detector,  args, out_path=out_path ,save_orig=args.save_orig)
+        except Exception as e:
+            print("Error:", e, video_path, audio_path)
+            import traceback
+            print(traceback.format_exc())
+
+        
+
+def main():
+    args = create_argparser().parse_args()
+    dist_util.setup_dist()
+    logger.configure(dir=args.sample_path, format_strs=["stdout", "log"])
+
+    logger.log("creating model...")
+    model, diffusion = tfg_create_model_and_diffusion(
+            **args_to_dict(args, tfg_model_and_diffusion_defaults().keys())
+    )
+    model.load_state_dict(
+            dist_util.load_state_dict(args.model_path, map_location='cpu')
+    )
+    model.to(dist_util.dev())
+    if args.use_fp16:
+        model.convert_to_fp16()
+    model.eval()
+
+    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, device='cuda' if torch.cuda.is_available() else 'cpu')
+
+    if args.generate_from_filelist:
+        generate_from_filelist(args.test_video_dir, args.filelist, model, diffusion, detector,  args)
+    else:
+        generate(args.video_path, args.audio_path, model, diffusion, detector,  args, out_path=args.out_path, save_orig=args.save_orig)
+
+
+def create_argparser():
+    defaults = dict(
+        # generate from a single audio-video pair
+        generate_from_filelist = False,
+        video_path = "",
+        audio_path = "",
+        out_path = None, 
+        save_orig = True,
+
+        #generate from filelist : generate_from_filelist = True
+        test_video_dir = "test_videos",
+        filelist = "test_filelist.txt",
+
+
+        use_fp16 = True,
+        #tfg specific
+        face_hide_percentage=0.5,
+        use_ref=False,
+        use_audio=False,
+        audio_as_style=False,
+        audio_as_style_encoder_mlp=False,
+        
+        #data args
+        nframes=1,
+        nrefer=0,
+        image_size=128,
+        syncnet_T = 5,
+        syncnet_mel_step_size = 16,
+        audio_frames_per_video = 16, #for tfg model, we use sound corresponding to 5 frames centred at that frame
+        audio_dim=80,
+        is_voxceleb2=True,
+
+        video_fps=25,
+        sample_rate=16000, #audio sampling rate
+        mel_steps_per_sec=80.,
+
+        #sampling args
+        clip_denoised=True, # not used in training
+        sampling_batch_size=2,
+        use_ddim=False,
+        model_path="",
+        sample_path="d2l_gen",
+        sample_partition="",
+        sampling_seed=None,
+        sampling_use_gt_for_ref=False,
+        sampling_ref_type='gt', #one of ['gt', 'first_frame', 'random']
+        sampling_input_type='gt', #one of ['gt', 'first_frame']
+        
+        # face detection args
+        face_det_batch_size=64,
+        pads = "0,0,0,0"
+    )
+    defaults.update(tfg_model_and_diffusion_defaults())
+    parser = argparse.ArgumentParser()
+    add_dict_to_argparser(parser, defaults)
+    return parser
+
+
+if __name__=="__main__":
+    main()

guided-diffusion/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

guided-diffusion/guided_diffusion/__init__.py

@@ -0,0 +1,3 @@
+"""
+Codebase for "Improved Denoising Diffusion Probabilistic Models".
+"""

guided-diffusion/guided_diffusion/dist_util.py

@@ -0,0 +1,94 @@
+"""
+Helpers for distributed training.
+"""
+
+import io
+import os
+import socket
+
+import blobfile as bf
+from mpi4py import MPI
+import torch as th
+import torch.distributed as dist
+
+# Change this to reflect your cluster layout.
+# The GPU for a given rank is (rank % GPUS_PER_NODE).
+GPUS_PER_NODE = 8
+
+SETUP_RETRY_COUNT = 3
+
+
+def setup_dist():
+    """
+    Setup a distributed process group.
+    """
+    if dist.is_initialized():
+        return
+    print("MPI.COMM_WORLD.Get_rank()", MPI.COMM_WORLD.Get_rank())
+    os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}"
+    print('os.environ["CUDA_VISIBLE_DEVICES"]', os.environ["CUDA_VISIBLE_DEVICES"])
+    comm = MPI.COMM_WORLD
+    backend = "gloo" if not th.cuda.is_available() else "nccl"
+
+    if backend == "gloo":
+        hostname = "localhost"
+    else:
+        hostname = socket.gethostbyname(socket.getfqdn())
+    os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
+    os.environ["RANK"] = str(comm.rank)
+    os.environ["WORLD_SIZE"] = str(comm.size)
+
+    port = comm.bcast(_find_free_port(), root=0)
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group(backend=backend, init_method="env://")
+
+
+def dev():
+    """
+    Get the device to use for torch.distributed.
+    """
+    if th.cuda.is_available():
+        return th.device(f"cuda")
+    return th.device("cpu")
+
+
+def load_state_dict(path, **kwargs):
+    """
+    Load a PyTorch file without redundant fetches across MPI ranks.
+    """
+    chunk_size = 2 ** 30  # MPI has a relatively small size limit
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        with bf.BlobFile(path, "rb") as f:
+            data = f.read()
+        num_chunks = len(data) // chunk_size
+        if len(data) % chunk_size:
+            num_chunks += 1
+        MPI.COMM_WORLD.bcast(num_chunks)
+        for i in range(0, len(data), chunk_size):
+            MPI.COMM_WORLD.bcast(data[i : i + chunk_size])
+    else:
+        num_chunks = MPI.COMM_WORLD.bcast(None)
+        data = bytes()
+        for _ in range(num_chunks):
+            data += MPI.COMM_WORLD.bcast(None)
+
+    return th.load(io.BytesIO(data), **kwargs)
+
+
+def sync_params(params):
+    """
+    Synchronize a sequence of Tensors across ranks from rank 0.
+    """
+    for p in params:
+        with th.no_grad():
+            dist.broadcast(p, 0)
+
+
+def _find_free_port():
+    try:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+    finally:
+        s.close()

guided-diffusion/guided_diffusion/fp16_util.py

@@ -0,0 +1,237 @@
+"""
+Helpers to train with 16-bit precision.
+"""
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from . import logger
+
+INITIAL_LOG_LOSS_SCALE = 20.0
+
+
+def convert_module_to_f16(l):
+    """
+    Convert primitive modules to float16.
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.half()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.half()
+
+
+def convert_module_to_f32(l):
+    """
+    Convert primitive modules to float32, undoing convert_module_to_f16().
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.float()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.float()
+
+
+def make_master_params(param_groups_and_shapes):
+    """
+    Copy model parameters into a (differently-shaped) list of full-precision
+    parameters.
+    """
+    master_params = []
+    for param_group, shape in param_groups_and_shapes:
+        master_param = nn.Parameter(
+            _flatten_dense_tensors(
+                [param.detach().float() for (_, param) in param_group]
+            ).view(shape)
+        )
+        master_param.requires_grad = True
+        master_params.append(master_param)
+    return master_params
+
+
+def model_grads_to_master_grads(param_groups_and_shapes, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    for master_param, (param_group, shape) in zip(
+        master_params, param_groups_and_shapes
+    ):
+        master_param.grad = _flatten_dense_tensors(
+            [param_grad_or_zeros(param) for (_, param) in param_group]
+        ).view(shape)
+
+
+def master_params_to_model_params(param_groups_and_shapes, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    # Without copying to a list, if a generator is passed, this will
+    # silently not copy any parameters.
+    for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes):
+        for (_, param), unflat_master_param in zip(
+            param_group, unflatten_master_params(param_group, master_param.view(-1))
+        ):
+            param.detach().copy_(unflat_master_param)
+
+
+def unflatten_master_params(param_group, master_param):
+    return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group])
+
+
+def get_param_groups_and_shapes(named_model_params):
+    named_model_params = list(named_model_params)
+    scalar_vector_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim <= 1],
+        (-1),
+    )
+    matrix_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim > 1],
+        (1, -1),
+    )
+    return [scalar_vector_named_params, matrix_named_params]
+
+
+def master_params_to_state_dict(
+    model, param_groups_and_shapes, master_params, use_fp16
+):
+    if use_fp16:
+        state_dict = model.state_dict()
+        for master_param, (param_group, _) in zip(
+            master_params, param_groups_and_shapes
+        ):
+            for (name, _), unflat_master_param in zip(
+                param_group, unflatten_master_params(param_group, master_param.view(-1))
+            ):
+                assert name in state_dict
+                state_dict[name] = unflat_master_param
+    else:
+        state_dict = model.state_dict()
+        for i, (name, _value) in enumerate(model.named_parameters()):
+            assert name in state_dict
+            state_dict[name] = master_params[i]
+    return state_dict
+
+
+def state_dict_to_master_params(model, state_dict, use_fp16):
+    if use_fp16:
+        named_model_params = [
+            (name, state_dict[name]) for name, _ in model.named_parameters()
+        ]
+        param_groups_and_shapes = get_param_groups_and_shapes(named_model_params)
+        master_params = make_master_params(param_groups_and_shapes)
+    else:
+        master_params = [state_dict[name] for name, _ in model.named_parameters()]
+    return master_params
+
+
+def zero_master_grads(master_params):
+    for param in master_params:
+        param.grad = None
+
+
+def zero_grad(model_params):
+    for param in model_params:
+        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
+        if param.grad is not None:
+            param.grad.detach_()
+            param.grad.zero_()
+
+
+def param_grad_or_zeros(param):
+    if param.grad is not None:
+        return param.grad.data.detach()
+    else:
+        return th.zeros_like(param)
+
+
+class MixedPrecisionTrainer:
+    def __init__(
+        self,
+        *,
+        model,
+        use_fp16=False,
+        fp16_scale_growth=1e-3,
+        initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE,
+    ):
+        self.model = model
+        self.use_fp16 = use_fp16
+        self.fp16_scale_growth = fp16_scale_growth
+
+        self.model_params = list(self.model.parameters())
+        self.master_params = self.model_params
+        self.param_groups_and_shapes = None
+        self.lg_loss_scale = initial_lg_loss_scale
+
+        if self.use_fp16:
+            self.param_groups_and_shapes = get_param_groups_and_shapes(
+                self.model.named_parameters()
+            )
+            self.master_params = make_master_params(self.param_groups_and_shapes)
+            self.model.convert_to_fp16()
+
+    def zero_grad(self):
+        zero_grad(self.model_params)
+
+    def backward(self, loss: th.Tensor):
+        if self.use_fp16:
+            loss_scale = 2 ** self.lg_loss_scale
+            (loss * loss_scale).backward()
+        else:
+            loss.backward()
+
+    def optimize(self, opt: th.optim.Optimizer):
+        if self.use_fp16:
+            return self._optimize_fp16(opt)
+        else:
+            return self._optimize_normal(opt)
+
+    def _optimize_fp16(self, opt: th.optim.Optimizer):
+        logger.logkv_mean("lg_loss_scale", self.lg_loss_scale)
+        model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params)
+        grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale)
+        if check_overflow(grad_norm):
+            self.lg_loss_scale -= 1
+            logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}")
+            zero_master_grads(self.master_params)
+            return False
+
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+
+        for p in self.master_params:
+            p.grad.mul_(1.0 / (2 ** self.lg_loss_scale))
+        opt.step()
+        zero_master_grads(self.master_params)
+        master_params_to_model_params(self.param_groups_and_shapes, self.master_params)
+        self.lg_loss_scale += self.fp16_scale_growth
+        return True
+
+    def _optimize_normal(self, opt: th.optim.Optimizer):
+        grad_norm, param_norm = self._compute_norms()
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+        opt.step()
+        return True
+
+    def _compute_norms(self, grad_scale=1.0):
+        grad_norm = 0.0
+        param_norm = 0.0
+        for p in self.master_params:
+            with th.no_grad():
+                param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2
+                if p.grad is not None:
+                    grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2
+        return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm)
+
+    def master_params_to_state_dict(self, master_params):
+        return master_params_to_state_dict(
+            self.model, self.param_groups_and_shapes, master_params, self.use_fp16
+        )
+
+    def state_dict_to_master_params(self, state_dict):
+        return state_dict_to_master_params(self.model, state_dict, self.use_fp16)
+
+
+def check_overflow(value):
+    return (value == float("inf")) or (value == -float("inf")) or (value != value)

guided-diffusion/guided_diffusion/gaussian_diffusion.py

@@ -0,0 +1,843 @@
+"""
+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.
+"""
+
+import enum
+import math
+
+import numpy as np
+import torch as th
+import os
+
+from . import dist_util
+from .nn import mean_flat
+from .losses import normal_kl, discretized_gaussian_log_likelihood
+
+
+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,
+        )
+    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)
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.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.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+class GaussianDiffusion:
+    """
+    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).
+    :param loss_variation: if True, then use composite loss
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type,
+        rescale_timesteps=False,
+        loss_variation=False,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.rescale_timesteps = rescale_timesteps
+        self.loss_variation = loss_variation
+
+        # Use float64 for accuracy.
+        betas = np.array(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 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 = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, self._scale_timesteps(t), **model_kwargs)
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == ModelVarType.LEARNED:
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(
+                    self.posterior_log_variance_clipped, t, x.shape
+                )
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            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 == ModelMeanType.PREVIOUS_X:
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
+            if self.model_mean_type == ModelMeanType.START_X:
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            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,
+        }
+
+    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_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 _scale_timesteps(self, t):
+        if self.rescale_timesteps:
+            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,
+        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,
+        shape,
+        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,
+        shape,
+        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
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            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)
+            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,
+        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,
+        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.
+        """
+        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
+        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_sample_loop(
+        self,
+        model,
+        shape,
+        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,
+        shape,
+        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
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            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)
+            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=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+        self, 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"]
+        )
+        if ("cond_img" in model_kwargs) and ("mask" in model_kwargs): #added by soumik
+            kl = kl*model_kwargs["mask"]
+        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
+        if ("cond_img" in model_kwargs) and ("mask" in model_kwargs): #added by soumik
+            decoder_nll=decoder_nll*model_kwargs["mask"]
+        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"]}
+
+
+    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, 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)

guided-diffusion/guided_diffusion/image_datasets.py

@@ -0,0 +1,167 @@
+import math
+import random
+
+from PIL import Image
+import blobfile as bf
+from mpi4py import MPI
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+
+
+def load_data(
+    *,
+    data_dir,
+    batch_size,
+    image_size,
+    class_cond=False,
+    deterministic=False,
+    random_crop=False,
+    random_flip=True,
+):
+    """
+    For a dataset, create a generator over (images, kwargs) pairs.
+
+    Each images is an NCHW float tensor, and the kwargs dict contains zero or
+    more keys, each of which map to a batched Tensor of their own.
+    The kwargs dict can be used for class labels, in which case the key is "y"
+    and the values are integer tensors of class labels.
+
+    :param data_dir: a dataset directory.
+    :param batch_size: the batch size of each returned pair.
+    :param image_size: the size to which images are resized.
+    :param class_cond: if True, include a "y" key in returned dicts for class
+                       label. If classes are not available and this is true, an
+                       exception will be raised.
+    :param deterministic: if True, yield results in a deterministic order.
+    :param random_crop: if True, randomly crop the images for augmentation.
+    :param random_flip: if True, randomly flip the images for augmentation.
+    """
+    if not data_dir:
+        raise ValueError("unspecified data directory")
+    all_files = _list_image_files_recursively(data_dir)
+    classes = None
+    if class_cond:
+        # Assume classes are the first part of the filename,
+        # before an underscore.
+        class_names = [bf.basename(path).split("_")[0] for path in all_files]
+        sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}
+        classes = [sorted_classes[x] for x in class_names]
+    dataset = ImageDataset(
+        image_size,
+        all_files,
+        classes=classes,
+        shard=MPI.COMM_WORLD.Get_rank(),
+        num_shards=MPI.COMM_WORLD.Get_size(),
+        random_crop=random_crop,
+        random_flip=random_flip,
+    )
+    if deterministic:
+        loader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True
+        )
+    else:
+        loader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True
+        )
+    while True:
+        yield from loader
+
+
+def _list_image_files_recursively(data_dir):
+    results = []
+    for entry in sorted(bf.listdir(data_dir)):
+        full_path = bf.join(data_dir, entry)
+        ext = entry.split(".")[-1]
+        if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]:
+            results.append(full_path)
+        elif bf.isdir(full_path):
+            results.extend(_list_image_files_recursively(full_path))
+    return results
+
+
+class ImageDataset(Dataset):
+    def __init__(
+        self,
+        resolution,
+        image_paths,
+        classes=None,
+        shard=0,
+        num_shards=1,
+        random_crop=False,
+        random_flip=True,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.local_images = image_paths[shard:][::num_shards]
+        self.local_classes = None if classes is None else classes[shard:][::num_shards]
+        self.random_crop = random_crop
+        self.random_flip = random_flip
+
+    def __len__(self):
+        return len(self.local_images)
+
+    def __getitem__(self, idx):
+        path = self.local_images[idx]
+        with bf.BlobFile(path, "rb") as f:
+            pil_image = Image.open(f)
+            pil_image.load()
+        pil_image = pil_image.convert("RGB")
+
+        if self.random_crop:
+            arr = random_crop_arr(pil_image, self.resolution)
+        else:
+            arr = center_crop_arr(pil_image, self.resolution)
+
+        if self.random_flip and random.random() < 0.5:
+            arr = arr[:, ::-1]
+
+        arr = arr.astype(np.float32) / 127.5 - 1
+
+        out_dict = {}
+        if self.local_classes is not None:
+            out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
+        return np.transpose(arr, [2, 0, 1]), out_dict
+
+
+def center_crop_arr(pil_image, image_size):
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
+
+
+def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
+    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
+    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
+    smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)
+
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * smaller_dim_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = smaller_dim_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = random.randrange(arr.shape[0] - image_size + 1)
+    crop_x = random.randrange(arr.shape[1] - image_size + 1)
+    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]

guided-diffusion/guided_diffusion/logger.py

@@ -0,0 +1,491 @@
+"""
+Logger copied from OpenAI baselines to avoid extra RL-based dependencies:
+https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py
+"""
+
+import os
+import sys
+import shutil
+import os.path as osp
+import json
+import time
+import datetime
+import tempfile
+import warnings
+from collections import defaultdict
+from contextlib import contextmanager
+from torch.utils.tensorboard import SummaryWriter
+
+DEBUG = 10
+INFO = 20
+WARN = 30
+ERROR = 40
+
+DISABLED = 50
+
+
+class KVWriter(object):
+    def writekvs(self, kvs):
+        raise NotImplementedError
+
+
+class SeqWriter(object):
+    def writeseq(self, seq):
+        raise NotImplementedError
+
+
+class HumanOutputFormat(KVWriter, SeqWriter):
+    def __init__(self, filename_or_file):
+        if isinstance(filename_or_file, str):
+            self.file = open(filename_or_file, "wt")
+            self.own_file = True
+        else:
+            assert hasattr(filename_or_file, "read"), (
+                "expected file or str, got %s" % filename_or_file
+            )
+            self.file = filename_or_file
+            self.own_file = False
+
+    def writekvs(self, kvs):
+        # Create strings for printing
+        key2str = {}
+        for (key, val) in sorted(kvs.items()):
+            if hasattr(val, "__float__"):
+                valstr = "%-8.3g" % val
+            else:
+                valstr = str(val)
+            key2str[self._truncate(key)] = self._truncate(valstr)
+
+        # Find max widths
+        if len(key2str) == 0:
+            print("WARNING: tried to write empty key-value dict")
+            return
+        else:
+            keywidth = max(map(len, key2str.keys()))
+            valwidth = max(map(len, key2str.values()))
+
+        # Write out the data
+        dashes = "-" * (keywidth + valwidth + 7)
+        lines = [dashes]
+        for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()):
+            lines.append(
+                "| %s%s | %s%s |"
+                % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val)))
+            )
+        lines.append(dashes)
+        self.file.write("\n".join(lines) + "\n")
+
+        # Flush the output to the file
+        self.file.flush()
+
+    def _truncate(self, s):
+        maxlen = 30
+        return s[: maxlen - 3] + "..." if len(s) > maxlen else s
+
+    def writeseq(self, seq):
+        seq = list(seq)
+        for (i, elem) in enumerate(seq):
+            self.file.write(elem)
+            if i < len(seq) - 1:  # add space unless this is the last one
+                self.file.write(" ")
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        if self.own_file:
+            self.file.close()
+
+
+class JSONOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "wt")
+
+    def writekvs(self, kvs):
+        for k, v in sorted(kvs.items()):
+            if hasattr(v, "dtype"):
+                kvs[k] = float(v)
+        self.file.write(json.dumps(kvs) + "\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class CSVOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "w+t")
+        self.keys = []
+        self.sep = ","
+
+    def writekvs(self, kvs):
+        # Add our current row to the history
+        extra_keys = list(kvs.keys() - self.keys)
+        extra_keys.sort()
+        if extra_keys:
+            self.keys.extend(extra_keys)
+            self.file.seek(0)
+            lines = self.file.readlines()
+            self.file.seek(0)
+            for (i, k) in enumerate(self.keys):
+                if i > 0:
+                    self.file.write(",")
+                self.file.write(k)
+            self.file.write("\n")
+            for line in lines[1:]:
+                self.file.write(line[:-1])
+                self.file.write(self.sep * len(extra_keys))
+                self.file.write("\n")
+        for (i, k) in enumerate(self.keys):
+            if i > 0:
+                self.file.write(",")
+            v = kvs.get(k)
+            if v is not None:
+                self.file.write(str(v))
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class TensorBoardOutputFormat(KVWriter):
+    """
+    Dumps key/value pairs into TensorBoard's numeric format.
+    """
+
+    def __init__(self, dir):
+        os.makedirs(dir, exist_ok=True)
+        self.dir = dir
+        self.step = -1
+        self.writer = SummaryWriter(self.dir)
+
+    def writekvs(self, kvs):
+        self.step = int(kvs["step"])
+        for k, v in sorted(kvs.items()):
+            self.writer.add_scalar(k, float(v), self.step)
+        self.writer.flush()
+    
+    def writeimage(self, key, image_tensor):
+        self.writer.add_image(key, image_tensor, self.step)
+        self.writer.flush()
+
+    def close(self):
+        if self.writer:
+            self.writer.close()
+            self.writer = None
+
+
+def make_output_format(format, ev_dir, log_suffix=""):
+    os.makedirs(ev_dir, exist_ok=True)
+    if format == "stdout":
+        return HumanOutputFormat(sys.stdout)
+    elif format == "log":
+        return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix))
+    elif format == "json":
+        return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix))
+    elif format == "csv":
+        return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix))
+    elif format == "tensorboard":
+        return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix))
+    else:
+        raise ValueError("Unknown format specified: %s" % (format,))
+
+
+# ================================================================
+# API
+# ================================================================
+
+def logimage(key, image_tensor):
+    """
+    Log one image to tensorboard
+    """
+    for fmt in get_current().output_formats:
+        if isinstance(fmt, TensorBoardOutputFormat):
+            tb_logger = fmt
+            tb_logger.writeimage(key, image_tensor)
+
+
+def logkv(key, val):
+    """
+    Log a value of some diagnostic
+    Call this once for each diagnostic quantity, each iteration
+    If called many times, last value will be used.
+    """
+    get_current().logkv(key, val)
+
+
+def logkv_mean(key, val):
+    """
+    The same as logkv(), but if called many times, values averaged.
+    """
+    get_current().logkv_mean(key, val)
+
+
+def logkvs(d):
+    """
+    Log a dictionary of key-value pairs
+    """
+    for (k, v) in d.items():
+        logkv(k, v)
+
+
+def dumpkvs():
+    """
+    Write all of the diagnostics from the current iteration
+    """
+    return get_current().dumpkvs()
+
+
+def getkvs():
+    return get_current().name2val
+
+
+def log(*args, level=INFO):
+    """
+    Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).
+    """
+    get_current().log(*args, level=level)
+
+
+def debug(*args):
+    log(*args, level=DEBUG)
+
+
+def info(*args):
+    log(*args, level=INFO)
+
+
+def warn(*args):
+    log(*args, level=WARN)
+
+
+def error(*args):
+    log(*args, level=ERROR)
+
+
+def set_level(level):
+    """
+    Set logging threshold on current logger.
+    """
+    get_current().set_level(level)
+
+
+def set_comm(comm):
+    get_current().set_comm(comm)
+
+
+def get_dir():
+    """
+    Get directory that log files are being written to.
+    will be None if there is no output directory (i.e., if you didn't call start)
+    """
+    return get_current().get_dir()
+
+
+record_tabular = logkv
+dump_tabular = dumpkvs
+
+
+@contextmanager
+def profile_kv(scopename):
+    logkey = "wait_" + scopename
+    tstart = time.time()
+    try:
+        yield
+    finally:
+        get_current().name2val[logkey] += time.time() - tstart
+
+
+def profile(n):
+    """
+    Usage:
+    @profile("my_func")
+    def my_func(): code
+    """
+
+    def decorator_with_name(func):
+        def func_wrapper(*args, **kwargs):
+            with profile_kv(n):
+                return func(*args, **kwargs)
+
+        return func_wrapper
+
+    return decorator_with_name
+
+
+# ================================================================
+# Backend
+# ================================================================
+
+
+def get_current():
+    if Logger.CURRENT is None:
+        _configure_default_logger()
+
+    return Logger.CURRENT
+
+
+class Logger(object):
+    DEFAULT = None  # A logger with no output files. (See right below class definition)
+    # So that you can still log to the terminal without setting up any output files
+    CURRENT = None  # Current logger being used by the free functions above
+
+    def __init__(self, dir, output_formats, comm=None):
+        self.name2val = defaultdict(float)  # values this iteration
+        self.name2cnt = defaultdict(int)
+        self.level = INFO
+        self.dir = dir
+        self.output_formats = output_formats
+        self.comm = comm
+
+    # Logging API, forwarded
+    # ----------------------------------------
+    def logkv(self, key, val):
+        self.name2val[key] = val
+
+    def logkv_mean(self, key, val):
+        oldval, cnt = self.name2val[key], self.name2cnt[key]
+        self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)
+        self.name2cnt[key] = cnt + 1
+
+    def dumpkvs(self):
+        if self.comm is None:
+            d = self.name2val
+        else:
+            d = mpi_weighted_mean(
+                self.comm,
+                {
+                    name: (val, self.name2cnt.get(name, 1))
+                    for (name, val) in self.name2val.items()
+                },
+            )
+            if self.comm.rank != 0:
+                d["dummy"] = 1  # so we don't get a warning about empty dict
+        out = d.copy()  # Return the dict for unit testing purposes
+        for fmt in self.output_formats:
+            if isinstance(fmt, KVWriter):
+                fmt.writekvs(d)
+        self.name2val.clear()
+        self.name2cnt.clear()
+        return out
+
+    def log(self, *args, level=INFO):
+        if self.level <= level:
+            self._do_log(args)
+
+    # Configuration
+    # ----------------------------------------
+    def set_level(self, level):
+        self.level = level
+
+    def set_comm(self, comm):
+        self.comm = comm
+
+    def get_dir(self):
+        return self.dir
+
+    def close(self):
+        for fmt in self.output_formats:
+            fmt.close()
+
+    # Misc
+    # ----------------------------------------
+    def _do_log(self, args):
+        for fmt in self.output_formats:
+            if isinstance(fmt, SeqWriter):
+                fmt.writeseq(map(str, args))
+
+
+def get_rank_without_mpi_import():
+    # check environment variables here instead of importing mpi4py
+    # to avoid calling MPI_Init() when this module is imported
+    for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]:
+        if varname in os.environ:
+            return int(os.environ[varname])
+    return 0
+
+
+def mpi_weighted_mean(comm, local_name2valcount):
+    """
+    Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110
+    Perform a weighted average over dicts that are each on a different node
+    Input: local_name2valcount: dict mapping key -> (value, count)
+    Returns: key -> mean
+    """
+    all_name2valcount = comm.gather(local_name2valcount)
+    if comm.rank == 0:
+        name2sum = defaultdict(float)
+        name2count = defaultdict(float)
+        for n2vc in all_name2valcount:
+            for (name, (val, count)) in n2vc.items():
+                try:
+                    val = float(val)
+                except ValueError:
+                    if comm.rank == 0:
+                        warnings.warn(
+                            "WARNING: tried to compute mean on non-float {}={}".format(
+                                name, val
+                            )
+                        )
+                else:
+                    name2sum[name] += val * count
+                    name2count[name] += count
+        return {name: name2sum[name] / name2count[name] for name in name2sum}
+    else:
+        return {}
+
+
+def configure(dir=None, format_strs=None, comm=None, log_suffix=""):
+    """
+    If comm is provided, average all numerical stats across that comm
+    """
+    if dir is None:
+        dir = os.getenv("OPENAI_LOGDIR")
+    if dir is None:
+        dir = osp.join(
+            tempfile.gettempdir(),
+            datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"),
+        )
+    assert isinstance(dir, str)
+    dir = os.path.expanduser(dir)
+    os.makedirs(os.path.expanduser(dir), exist_ok=True)
+
+    rank = get_rank_without_mpi_import()
+    if rank > 0:
+        log_suffix = log_suffix + "-rank%03i" % rank
+
+    if format_strs is None:
+        if rank == 0:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv,tensorboard").split(",")
+        else:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",")
+    format_strs = filter(None, format_strs)
+    output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs]
+
+    Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm)
+    if output_formats:
+        log("Logging to %s" % dir)
+
+
+def _configure_default_logger():
+    configure()
+    Logger.DEFAULT = Logger.CURRENT
+
+
+def reset():
+    if Logger.CURRENT is not Logger.DEFAULT:
+        Logger.CURRENT.close()
+        Logger.CURRENT = Logger.DEFAULT
+        log("Reset logger")
+
+
+@contextmanager
+def scoped_configure(dir=None, format_strs=None, comm=None):
+    prevlogger = Logger.CURRENT
+    configure(dir=dir, format_strs=format_strs, comm=comm)
+    try:
+        yield
+    finally:
+        Logger.CURRENT.close()
+        Logger.CURRENT = prevlogger
+

guided-diffusion/guided_diffusion/losses.py

@@ -0,0 +1,77 @@
+"""
+Helpers for various likelihood-based losses. These are ported from the original
+Ho et al. diffusion models codebase:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py
+"""
+
+import numpy as np
+
+import torch as th
+
+
+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

guided-diffusion/guided_diffusion/lpips.py

@@ -0,0 +1,20 @@
+from lpips_pytorch import LPIPS
+import torch
+
+class LPIPS1(LPIPS):
+    r"""
+    Overrriding the LPIPS to send loss without reducing the batch
+    Arguments:
+        net_type (str): the network type to compare the features: 
+                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.
+        version (str): the version of LPIPS. Default: 0.1.
+    """
+    def __init__(self, net_type: str = 'alex', version: str = '0.1'):
+        super(LPIPS1, self).__init__(net_type = 'alex', version ='0.1')
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        feat_x, feat_y = self.net(x), self.net(y)
+        diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]
+        res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]
+        # return torch.sum(torch.cat(res, 0), 0, True)
+        return torch.sum(torch.cat(res, 1), 1, True)

guided-diffusion/guided_diffusion/nn.py

@@ -0,0 +1,170 @@
+"""
+Various utilities for neural networks.
+"""
+
+import math
+
+import torch as th
+import torch.nn as nn
+
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+    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(32, 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 checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+
+class CheckpointFunction(th.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        with th.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with th.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = th.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads

guided-diffusion/guided_diffusion/resample.py

@@ -0,0 +1,154 @@
+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)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(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, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()

guided-diffusion/guided_diffusion/respace.py

@@ -0,0 +1,128 @@
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+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)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    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, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # 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:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, 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):
+        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:
+    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 __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        if self.rescale_timesteps:
+            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)

guided-diffusion/guided_diffusion/script_util.py

@@ -0,0 +1,614 @@
+import argparse
+import inspect
+
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+from .unet import SuperResModel, UNetModel, EncoderUNetModel, TFGModel
+
+NUM_CLASSES = 1000
+
+
+def diffusion_defaults():
+    """
+    Defaults for image and classifier training.
+    """
+    return dict(
+        learn_sigma=False,
+        diffusion_steps=1000,
+        noise_schedule="linear",
+        timestep_respacing="",
+        use_kl=False,
+        predict_xstart=False,
+        rescale_timesteps=False,
+        rescale_learned_sigmas=False,
+        loss_variation=0, #added by soumik
+    )
+
+
+def classifier_defaults():
+    """
+    Defaults for classifier models.
+    """
+    return dict(
+        image_size=64,
+        classifier_use_fp16=False,
+        classifier_width=128,
+        classifier_depth=2,
+        classifier_attention_resolutions="32,16,8",  # 16
+        classifier_use_scale_shift_norm=True,  # False
+        classifier_resblock_updown=True,  # False
+        classifier_pool="attention",
+    )
+
+
+def model_and_diffusion_defaults():
+    """
+    Defaults for image training.
+    """
+    res = dict(
+        image_size=64,
+        num_channels=128,
+        num_res_blocks=2,
+        num_heads=4,
+        num_heads_upsample=-1,
+        num_head_channels=-1,
+        attention_resolutions="16,8",
+        channel_mult="",
+        dropout=0.0,
+        class_cond=False,
+        use_checkpoint=False,
+        use_scale_shift_norm=True,
+        resblock_updown=False,
+        use_fp16=False,
+        use_new_attention_order=False,
+    )
+    res.update(diffusion_defaults())
+    return res
+
+
+def classifier_and_diffusion_defaults():
+    res = classifier_defaults()
+    res.update(diffusion_defaults())
+    return res
+
+
+def create_model_and_diffusion(
+    image_size,
+    class_cond,
+    learn_sigma,
+    num_channels,
+    num_res_blocks,
+    channel_mult,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    attention_resolutions,
+    dropout,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    use_checkpoint,
+    use_scale_shift_norm,
+    resblock_updown,
+    use_fp16,
+    use_new_attention_order,
+):
+    model = create_model(
+        image_size,
+        num_channels,
+        num_res_blocks,
+        channel_mult=channel_mult,
+        learn_sigma=learn_sigma,
+        class_cond=class_cond,
+        use_checkpoint=use_checkpoint,
+        attention_resolutions=attention_resolutions,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        dropout=dropout,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+        use_new_attention_order=use_new_attention_order,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return model, diffusion
+
+
+def create_model(
+    image_size,
+    num_channels,
+    num_res_blocks,
+    channel_mult="",
+    learn_sigma=False,
+    class_cond=False,
+    use_checkpoint=False,
+    attention_resolutions="16",
+    num_heads=1,
+    num_head_channels=-1,
+    num_heads_upsample=-1,
+    use_scale_shift_norm=False,
+    dropout=0,
+    resblock_updown=False,
+    use_fp16=False,
+    use_new_attention_order=False,
+):
+    if channel_mult == "":
+        if image_size == 512:
+            channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
+        elif image_size == 256:
+            channel_mult = (1, 1, 2, 2, 4, 4)
+        elif image_size == 128:
+            channel_mult = (1, 1, 2, 3, 4)
+        elif image_size == 64:
+            channel_mult = (1, 2, 3, 4)
+        else:
+            raise ValueError(f"unsupported image size: {image_size}")
+    else:
+        channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(","))
+
+    attention_ds = []
+    for res in attention_resolutions.split(","):
+        attention_ds.append(image_size // int(res))
+
+    return UNetModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=dropout,
+        channel_mult=channel_mult,
+        num_classes=(NUM_CLASSES if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        use_fp16=use_fp16,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_new_attention_order=use_new_attention_order,
+    )
+
+
+def create_classifier_and_diffusion(
+    image_size,
+    classifier_use_fp16,
+    classifier_width,
+    classifier_depth,
+    classifier_attention_resolutions,
+    classifier_use_scale_shift_norm,
+    classifier_resblock_updown,
+    classifier_pool,
+    learn_sigma,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+):
+    classifier = create_classifier(
+        image_size,
+        classifier_use_fp16,
+        classifier_width,
+        classifier_depth,
+        classifier_attention_resolutions,
+        classifier_use_scale_shift_norm,
+        classifier_resblock_updown,
+        classifier_pool,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return classifier, diffusion
+
+
+def create_classifier(
+    image_size,
+    classifier_use_fp16,
+    classifier_width,
+    classifier_depth,
+    classifier_attention_resolutions,
+    classifier_use_scale_shift_norm,
+    classifier_resblock_updown,
+    classifier_pool,
+):
+    if image_size == 512:
+        channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
+    elif image_size == 256:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif image_size == 128:
+        channel_mult = (1, 1, 2, 3, 4)
+    elif image_size == 64:
+        channel_mult = (1, 2, 3, 4)
+    else:
+        raise ValueError(f"unsupported image size: {image_size}")
+
+    attention_ds = []
+    for res in classifier_attention_resolutions.split(","):
+        attention_ds.append(image_size // int(res))
+
+    return EncoderUNetModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=classifier_width,
+        out_channels=1000,
+        num_res_blocks=classifier_depth,
+        attention_resolutions=tuple(attention_ds),
+        channel_mult=channel_mult,
+        use_fp16=classifier_use_fp16,
+        num_head_channels=64,
+        use_scale_shift_norm=classifier_use_scale_shift_norm,
+        resblock_updown=classifier_resblock_updown,
+        pool=classifier_pool,
+    )
+
+
+def sr_model_and_diffusion_defaults():
+    res = model_and_diffusion_defaults()
+    res["large_size"] = 256
+    res["small_size"] = 64
+    arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0]
+    for k in res.copy().keys():
+        if k not in arg_names:
+            del res[k]
+    return res
+
+
+def sr_create_model_and_diffusion(
+    large_size,
+    small_size,
+    class_cond,
+    learn_sigma,
+    num_channels,
+    num_res_blocks,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    attention_resolutions,
+    dropout,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    use_checkpoint,
+    use_scale_shift_norm,
+    resblock_updown,
+    use_fp16,
+):
+    model = sr_create_model(
+        large_size,
+        small_size,
+        num_channels,
+        num_res_blocks,
+        learn_sigma=learn_sigma,
+        class_cond=class_cond,
+        use_checkpoint=use_checkpoint,
+        attention_resolutions=attention_resolutions,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        dropout=dropout,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return model, diffusion
+
+
+def sr_create_model(
+    large_size,
+    small_size,
+    num_channels,
+    num_res_blocks,
+    learn_sigma,
+    class_cond,
+    use_checkpoint,
+    attention_resolutions,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    use_scale_shift_norm,
+    dropout,
+    resblock_updown,
+    use_fp16,
+):
+    _ = small_size  # hack to prevent unused variable
+
+    if large_size == 512:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif large_size == 256:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif large_size == 64:
+        channel_mult = (1, 2, 3, 4)
+    else:
+        raise ValueError(f"unsupported large size: {large_size}")
+
+    attention_ds = []
+    for res in attention_resolutions.split(","):
+        attention_ds.append(large_size // int(res))
+
+    return SuperResModel(
+        image_size=large_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=dropout,
+        channel_mult=channel_mult,
+        num_classes=(NUM_CLASSES if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+    )
+
+
+def create_gaussian_diffusion(
+    *,
+    steps=1000,
+    learn_sigma=False,
+    sigma_small=False,
+    noise_schedule="linear",
+    use_kl=False,
+    predict_xstart=False,
+    rescale_timesteps=False,
+    rescale_learned_sigmas=False,
+    timestep_respacing="",
+    loss_variation=0,
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, steps)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if not timestep_respacing:
+        timestep_respacing = [steps]
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=(
+            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
+        ),
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type,
+        rescale_timesteps=rescale_timesteps,
+        loss_variation=loss_variation, #added by soumik
+    )
+
+
+def add_dict_to_argparser(parser, default_dict):
+    for k, v in default_dict.items():
+        v_type = type(v)
+        if v is None:
+            v_type = str
+        elif isinstance(v, bool):
+            v_type = str2bool
+        parser.add_argument(f"--{k}", default=v, type=v_type)
+
+
+def args_to_dict(args, keys):
+    return {k: getattr(args, k) for k in keys}
+
+
+def str2bool(v):
+    """
+    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("boolean value expected")
+
+#________________________________ tfg model ________________________________#
+def tfg_model_and_diffusion_defaults():
+    res = model_and_diffusion_defaults()
+    arg_names = inspect.getfullargspec(tfg_create_model_and_diffusion)[0]
+    for k in res.copy().keys():
+        if k not in arg_names:
+            del res[k]
+    
+    #tfg args
+    res["use_ref"]=False
+    res["nframes"]=1
+    res["nrefer"]=0
+    res["use_audio"]=False
+    res["audio_encoder_kwargs"]={}
+    res["audio_as_style"]=False
+    res["audio_as_style_encoder_mlp"]=False
+    return res
+    
+def tfg_create_model_and_diffusion(
+    image_size,
+    class_cond,
+    learn_sigma, 
+    num_channels,
+    num_res_blocks,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    attention_resolutions,
+    dropout,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing, 
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    use_checkpoint,
+    use_scale_shift_norm,
+    resblock_updown,
+    use_fp16,
+    use_ref,
+    nframes,
+    nrefer,
+    use_audio,
+    audio_encoder_kwargs,
+    audio_as_style,
+    audio_as_style_encoder_mlp,
+    loss_variation,
+):
+    model = tfg_create_model(
+        image_size,
+        num_channels,
+        num_res_blocks,
+        learn_sigma=learn_sigma,
+        class_cond=class_cond,
+        use_checkpoint=use_checkpoint,
+        attention_resolutions=attention_resolutions,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        dropout=dropout,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+        use_ref=use_ref,
+        nframes=nframes,
+        nrefer=nrefer,
+        use_audio=use_audio,
+        audio_encoder_kwargs=audio_encoder_kwargs,
+        audio_as_style=audio_as_style,
+        audio_as_style_encoder_mlp=audio_as_style_encoder_mlp,
+    )
+
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+        loss_variation=loss_variation,
+    )
+    return model, diffusion
+
+def tfg_create_model(
+    image_size,
+    num_channels,
+    num_res_blocks,
+    learn_sigma,
+    class_cond,
+    use_checkpoint,
+    attention_resolutions,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    use_scale_shift_norm,
+    dropout,
+    resblock_updown,
+    use_fp16,
+    use_ref,
+    nframes,
+    nrefer,
+    use_audio,
+    audio_encoder_kwargs,
+    audio_as_style,
+    audio_as_style_encoder_mlp,
+):
+    
+    if image_size == 512:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif image_size == 256:
+        channel_mult = (1, 1, 2, 3, 4, 4)
+    elif image_size ==  128:
+        channel_mult = (1, 1, 2, 3, 4)
+    elif image_size == 64:
+        channel_mult = (1, 2, 3, 4)
+    else:
+        raise ValueError(f"unsupported large size: {image_size}")
+
+    attention_ds = []
+    if "-1" not in attention_resolutions: # -1 = no attention
+        for res in attention_resolutions.split(","):
+            attention_ds.append(image_size // int(res))
+    
+    return TFGModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=dropout,
+        channel_mult=channel_mult,
+        num_classes=(NUM_CLASSES if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        use_fp16=use_fp16,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_ref=use_ref,
+        nframes=nframes,
+        nrefer=nrefer,
+        use_audio=use_audio,
+        audio_encoder_kwargs=audio_encoder_kwargs,
+        audio_as_style=audio_as_style,
+        audio_as_style_encoder_mlp=audio_as_style_encoder_mlp
+    )
+    
+
+
+
+
+
+        

guided-diffusion/guided_diffusion/tfg_data_util.py

@@ -0,0 +1,75 @@
+import torch
+
+def normalise2(tensor):
+    '''[0,1] -> [-1,1]'''
+    return (tensor*2 - 1.).clamp(-1,1)
+
+def tfg_data(dataloader, face_hide_percentage, use_ref, use_audio):#, sampling_use_gt_for_ref=False, noise = None):
+    def inf_gen(generator):
+        while True:
+            yield from generator
+    data = inf_gen(dataloader)
+    for batch in data:
+        img_batch, model_kwargs = tfg_process_batch(batch, face_hide_percentage, use_ref, use_audio)
+        yield img_batch, model_kwargs
+        
+
+def tfg_process_batch(batch, face_hide_percentage, use_ref=False, use_audio=False, sampling_use_gt_for_ref=False, noise = None):
+    model_kwargs = {}
+    B, F,C, H, W = batch["image"].shape
+    img_batch = normalise2(batch["image"].reshape(B*F, C, H, W).contiguous())
+    model_kwargs = tfg_add_cond_inputs(img_batch, model_kwargs, face_hide_percentage, noise)
+    if use_ref:
+        model_kwargs = tfg_add_reference(batch, model_kwargs, sampling_use_gt_for_ref)
+    if use_audio:
+        model_kwargs = tfg_add_audio(batch,model_kwargs)
+    return img_batch, model_kwargs
+
+def tfg_add_reference(batch, model_kwargs, sampling_use_gt_for_ref=False):
+    # assuming nrefer = 1
+    #[B, nframes, C, H, W] -> #[B*nframes, C, H, W]
+    if sampling_use_gt_for_ref:
+        B, F,C, H, W = batch["image"].shape
+        img_batch = normalise2(batch["image"].reshape(B*F, C, H, W).contiguous())
+        model_kwargs["ref_img"] = img_batch
+    else:
+        _, _, C, H , W =  batch["ref_img"].shape
+        ref_img = normalise2(batch["ref_img"].reshape(-1, C, H, W).contiguous())
+        model_kwargs["ref_img"] = ref_img
+    return model_kwargs
+
+def tfg_add_audio(batch, model_kwargs):
+    # unet needs [BF, h, w] as input
+    B, F, _, h, w = batch["indiv_mels"].shape
+    indiv_mels = batch["indiv_mels"] # [B, F, 1, h, w]
+    indiv_mels = indiv_mels.squeeze(dim=2).reshape(B*F, h , w)
+    model_kwargs["indiv_mels"] = indiv_mels
+    # syncloss needs [B, 1, 80, 16] as input
+    if "mel" in batch:
+        mel = batch["mel"] #[B, 1, h, w]
+        model_kwargs["mel"]=mel
+    return model_kwargs
+
+def tfg_add_cond_inputs(img_batch, model_kwargs, face_hide_percentage, noise=None):
+    B, C, H, W = img_batch.shape
+    mask = torch.zeros(B,1,H,W)
+    mask_start_idx = int (H*(1-face_hide_percentage))
+    mask[:,:,mask_start_idx:,:]=1.
+    if noise is None:
+        noise = torch.randn_like(img_batch)
+    assert noise.shape == img_batch.shape, "Noise shape != Image shape"
+    cond_img = img_batch *(1. - mask)+mask*noise
+
+    model_kwargs["cond_img"] = cond_img
+    model_kwargs["mask"] = mask
+    return model_kwargs
+
+
+def get_n_params(model):
+    pp=0
+    for p in list(model.parameters()):
+        nn=1
+        for s in list(p.size()):
+            nn=nn*s
+        pp+=nn
+    return pp

guided-diffusion/guided_diffusion/unet.py

@@ -0,0 +1,1275 @@
+from abc import abstractmethod
+
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .fp16_util import convert_module_to_f16, convert_module_to_f32
+from .nn import (
+    checkpoint,
+    conv_nd,
+    linear,
+    avg_pool_nd,
+    zero_module,
+    normalization,
+    timestep_embedding,
+)
+
+
+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]
+
+
+class TimestepBlock(nn.Module):
+    """
+    Any module where forward() takes timestep embeddings as a second argument.
+    """
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        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):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+        return x
+
+
+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, stride=None):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if stride is None:
+            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 ResBlock(TimestepBlock):
+    """
+    A residual block that can optionally change the number of channels.
+
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+        down_stride = None,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.dims = dims
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims, stride = down_stride)
+            self.x_upd = Downsample(channels, False, dims, stride = down_stride)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        return checkpoint(
+            self._forward, (x, emb), self.parameters(), self.use_checkpoint
+        )
+
+    def _forward(self, x, emb):
+        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)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = th.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+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 checkpoint(self._forward, (x,), self.parameters(), True)
+
+    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 UNetModel(nn.Module):
+    """
+    The full UNet model with attention and timestep embedding.
+
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.dims = dims
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.use_fp16 = use_fp16
+        self.dtype = th.float16 if use_fp16 else th.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.resblock_updown = resblock_updown
+
+        time_embed_dim = model_channels * 4
+        self.time_embed_dim = time_embed_dim
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if self.num_classes is not None:
+            self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+
+        ch = input_ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(model_channels * mult),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(model_channels * mult)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads_upsample,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    layers.append(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
+        )
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps, y=None):
+        """
+        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.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+
+        hs = []
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        if self.num_classes is not None:
+            assert y.shape == (x.shape[0],)
+            emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            hs.append(h)
+        h = self.middle_block(h, emb)
+        for module in self.output_blocks:
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(h, emb)
+        h = h.type(x.dtype)
+        return self.out(h)
+
+
+class SuperResModel(UNetModel):
+    """
+    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)
+
+
+class EncoderUNetModel(nn.Module):
+    """
+    The half UNet model with attention and timestep embedding.
+
+    For usage, see UNet.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        pool="adaptive",
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.dtype = th.float16 if use_fp16 else th.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+        self.pool = pool
+        if pool == "adaptive":
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                nn.AdaptiveAvgPool2d((1, 1)),
+                zero_module(conv_nd(dims, ch, out_channels, 1)),
+                nn.Flatten(),
+            )
+        elif pool == "attention":
+            assert num_head_channels != -1
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                AttentionPool2d(
+                    (image_size // ds), ch, num_head_channels, out_channels
+                ),
+            )
+        elif pool == "spatial":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                nn.ReLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        elif pool == "spatial_v2":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                normalization(2048),
+                nn.SiLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        else:
+            raise NotImplementedError(f"Unexpected {pool} pooling")
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps):
+        """
+        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.
+        """
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        results = []
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            if self.pool.startswith("spatial"):
+                results.append(h.type(x.dtype).mean(dim=(2, 3)))
+        h = self.middle_block(h, emb)
+        if self.pool.startswith("spatial"):
+            results.append(h.type(x.dtype).mean(dim=(2, 3)))
+            h = th.cat(results, axis=-1)
+            return self.out(h)
+        else:
+            h = h.type(x.dtype)
+            return self.out(h)
+
+
+#________________________________ tfg model ________________________________#
+class TFGModel(UNetModel):
+    '''
+    Talking Face Generation using UNet model
+    '''
+    def __init__(self, 
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        *args,
+        use_ref = False,
+        nframes = 1,
+        nrefer = 0,
+        use_audio = False,
+        audio_encoder_kwargs=None,
+        audio_as_style=False, # audio conditioned as style instead of concatenation in the middle
+        audio_as_style_encoder_mlp=False, # use mlp instead of audio encoder
+        **kwargs
+    ):
+        if use_ref:
+            super().__init__(image_size, in_channels * (1+1+nrefer), model_channels, out_channels * 1, *args, **kwargs)
+        else:
+            super().__init__(image_size, in_channels * (1+1), model_channels, out_channels * 1, *args, **kwargs)
+        
+
+        self.use_ref = use_ref
+        self.nframes = nframes
+        self.nrefer = nrefer
+        self.use_audio = use_audio
+
+        if self.use_audio:
+            if audio_encoder_kwargs is not None:
+                self.audio_encoder_kwargs = audio_encoder_kwargs
+            else:
+                self.audio_encoder_kwargs = {}
+
+            self.audio_as_style = audio_as_style
+            self.audio_as_style_encoder_mlp = audio_as_style_encoder_mlp
+            
+            self.audio_encoder = TFGAudioEncoder(
+                nframes = self.nframes,
+                dropout = self.dropout,
+                conv_resample = self.conv_resample,
+                dims = self.dims,
+                use_checkpoints = self.use_checkpoint,
+                use_fp16 = self.use_fp16,
+                use_scale_shift_norm = self.use_scale_shift_norm,
+                resblock_updown = self.resblock_updown,
+                **self.audio_encoder_kwargs
+            )
+
+            if not self.audio_as_style:
+                #concatenate audio encoding to the video encoding
+                old_middle_block_head = self.middle_block[0]
+                mid_img_ch = old_middle_block_head.channels
+                mid_aud_ch = self.audio_encoder.out_channels
+                self.middle_block[0] = ResBlock(
+                    mid_img_ch + mid_aud_ch, #combined image and audio channels
+                    old_middle_block_head.emb_channels,
+                    old_middle_block_head.dropout,
+                    out_channels = old_middle_block_head.out_channels,
+                    dims = old_middle_block_head.dims,
+                    use_checkpoint=old_middle_block_head.use_checkpoint,
+                    use_scale_shift_norm=old_middle_block_head.use_scale_shift_norm,
+                )
+            else: # audio as style
+                if self.audio_as_style_encoder_mlp:
+                    old_conv_encoder = self.audio_encoder
+                    audio_dim = old_conv_encoder.audio_dim 
+                    audio_frames_per_video = old_conv_encoder.audio_frames_per_video
+                    self.audio_encoder = nn.Sequential(
+                        nn.Flatten(),
+                        linear(audio_dim+audio_frames_per_video, self.time_embed_dim),
+                        normalization(self.time_embed_dim),
+                        nn.SiLU(),
+                        linear(self.time_embed_dim, self.time_embed_dim),
+                    )
+                else: # use conv_encoder+mlp to get style
+                    # similar to the classifier defined
+                    self.audio_encoder_to_style = nn.Sequential(
+                        normalization(self.audio_encoder.out_channels),
+                        nn.SiLU(),
+                        nn.AdaptiveAvgPool2d((1,1)),
+                        zero_module( #-> makes inital weights 0
+                            conv_nd(self.dims, self.audio_encoder.out_channels, self.time_embed_dim, 1)
+                        ),
+                        nn.Flatten(),
+                    )
+                
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+        if self.use_audio:
+            self.audio_encoder.apply(convert_module_to_f16)
+            if self.audio_as_style:
+                self.audio_encoder_to_style.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+        if self.use_audio:
+            self.audio_encoder.apply(convert_module_to_f32)
+            if self.audio_as_style:
+                self.audio_encoder_to_style.apply(convert_module_to_f32)
+ 
+    def forward(self, x, timesteps, cond_img=None, mask = None, ref_img=None, indiv_mels=None, **kwargs):
+
+        #preprocessing
+        x = x * mask + (1. - mask) * cond_img # mask the top half of the input
+        x = th.cat([x,cond_img], dim=1)
+        if self.use_ref:
+            x=th.cat([x,ref_img], dim=1)
+
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+
+        if self.use_audio:
+            if self.audio_as_style:
+                #audio encoder
+                if self.audio_as_style_encoder_mlp:#mlp uses fp32
+                    a = self.audio_encoder(indiv_mels)
+                    a = self.audio_encoder_to_style(a)
+                    a = a.type(self.dtype)
+                else:# conv uses fp16
+                    a = indiv_mels.type(self.dtype)
+                    a = self.audio_encoder(a)
+                    a= self.audio_encoder_to_style(a)
+                #combine
+                emb = emb + a
+                #video encoder
+                hs = []
+                h = x.type(self.dtype)
+                for module in self.input_blocks:
+                    h = module(h, emb)
+                    hs.append(h)
+
+            else: # concat audio in the middle
+                #audio encoder
+                a = indiv_mels.type(self.dtype)
+                a = self.audio_encoder(a)
+                #video encoder
+                hs = []
+                h = x.type(self.dtype)
+                for module in self.input_blocks:
+                    h = module(h, emb)
+                    hs.append(h)
+                #combine
+                h = th.cat([h,a], dim=1)
+        
+        #middle block
+        h = self.middle_block(h, emb)
+
+        # decoder
+        for module in self.output_blocks:
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(h, emb)
+        h = h.type(x.dtype)
+        return self.out(h)
+
+
+class TFGAudioEncoder(nn.Module):
+    """
+    Audio Encoder
+
+    with audio_dim = 80, 
+        audio_frames_per_video = 16
+        init_spatial_dim = 64
+        model_channels=32
+        channel_mult=(1,2,3,4)
+
+    following are the output shapes ->
+                init: [BF, 80, 16]
+                after in_block: [BF, 64, 16]
+                adding new dim [BF, 1, 64, 16]
+                encoder block before entering the loop: [BF, 32, 64, 16]
+                level: 0
+                0 _ 0 : [BF, 32, 64, 16]
+                0 _ 1 : [BF, 32, 64, 16]
+                0 _ 2 : [BF, 32, 32, 16]
+                level: 1
+                1 _ 0 : [BF, 64, 32, 16]
+                1 _ 1 : [BF, 64, 32, 16]
+                1 _ 2 : [BF, 64, 16, 16]
+                level: 2
+                2 _ 0 : [BF, 96, 16, 16]
+                2 _ 1 : [BF, 96, 16, 16]
+                2 _ 2 : [BF, 96, 8, 8]
+                level: 3
+                3 _ 0 : [BF, 128, 8, 8]
+                3 _ 1 : [BF, 128, 8, 8]
+                middle block: [BF, 128, 8, 8]
+                out: [BF, 128, 8, 8]
+    """
+    def __init__(
+        self, 
+        audio_dim = 80, 
+        audio_frames_per_video = 16, 
+        nframes=1,
+        
+        init_spatial_dim = 64,
+        model_channels=32,
+        out_channels=-1,
+        num_res_blocks=2,
+        dropout=0,
+        channel_mult=(1,2,3,4), #(1,1,2,4,8),
+        conv_resample=True,
+        dims=2,
+        use_checkpoint = False,
+        use_fp16=False,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        **kwargs
+    ):
+        super().__init__()
+        self.audio_dim = audio_dim
+        self.audio_frames_per_video = audio_frames_per_video
+        self.nframes = nframes
+        self.model_channels = model_channels
+        self.out_channels = out_channels if out_channels > 0 else model_channels * channel_mult[-1]
+        self.num_res_blocks = num_res_blocks
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.dtype = th.float16 if use_fp16 else th.float32
+        
+        time_embed_dim = model_channels
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        ch = int(channel_mult[0] * model_channels)
+        # init_spatial_dim = 4 * ( 2** (len(channel_mult)-1))
+
+        # convert spatial dim 80->64 using Conv1D: [N*F,80,16] -> [N*F, 64, 16]
+        _conv_dim, _in_channels, _out_channels = 1, self.audio_dim, init_spatial_dim
+        self.input_block = TimestepEmbedSequential(
+            conv_nd(_conv_dim, _in_channels, _out_channels, 3, padding=1),
+            normalization(_out_channels),
+            nn.SiLU()
+        )
+
+        # manually reshape [N*F, 64, 16] -> [N*F, 1,  64, 16] in __forward__()
+        
+        # [NF, 1, 64, 16] -> [N*F, model_channels*channel_mult[0], 64, 16] 
+        # can't use a resblock, bc groupnorm needs 32 sized group of channels
+        self.encoder_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(
+                    conv_nd(dims, 1, ch, 3, padding=1 )
+                )
+
+            ]
+        ) 
+
+        self._feature_size = ch
+        input_block_chans = [ch]
+
+        ds = 1 
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult*model_channels),
+                        dims = dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult*model_channels)
+                self.encoder_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult)-1:
+                out_ch = ch
+                self.encoder_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims = dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down = True,
+                            down_stride = (2,1) if (init_spatial_dim//ds) > self.audio_frames_per_video else (2,2),
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch,
+                            down_stride = (2,1) if (init_spatial_dim//ds) > self.audio_frames_per_video else (2,2),
+                            )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds*=2
+                self._feature_size += ch
+        
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                out_channels=self.out_channels,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+
+        # self.out = Upsample(self.out_channels, False, dims)
+        self.out = nn.Identity()
+        
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+
+    
+    def forward(self, x):
+        h = x.type(self.dtype)
+        BF_in, H_in, W_in = h.shape
+
+        #Fixed time embedding,(for using the same modules)
+        t = th.zeros(BF_in, dtype=th.long, device = x.device)
+        emb = self.time_embed(timestep_embedding(t, self.model_channels))
+
+        #80 -> 64 using Conv1D
+        h= self.input_block(h, emb)
+        _, H, W = h.shape
+        h = h.reshape(BF_in, 1, H, W) #[B*F, 64, 16] -> [B*F, 1, 64, 4]
+        # call encoder blocks
+        for module in self.encoder_blocks:
+            h = module(h,emb)
+        h = self.middle_block(h,emb)
+        h = h.type(x.dtype)
+        return self.out(h) # -> [B*F, 256, 8, 8]
+#______________________________________________________________________#
+
+            
+

guided-diffusion/setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup
+
+setup(
+    name="guided-diffusion",
+    py_modules=["guided_diffusion"],
+    install_requires=["blobfile>=1.0.5", "torch", "tqdm"],
+)

requirements.txt

@@ -0,0 +1,11 @@
+librosa==0.9.2
+opencv-python==4.5.5.64
+opencv-contrib-python==4.6.0.66
+tensorboard==2.11.0
+tqdm==4.64.1
+mpi4py-mpich==3.1.2
+av==9.2.0
+torch --extra-index-url https://download.pytorch.org/whl/cu113
+torchvision --extra-index-url https://download.pytorch.org/whl/cu113
+torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
+-e ./guided-diffusion

scripts/inference.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+#set paths and arguments
+real_video_root='dataset/VoxCeleb2/vox2_test_mp4/mp4/'
+model_path="checkpoints/checkpoint.pt"
+sample_path="output_dir"
+sample_mode="cross" # or "reconstruction"
+NUM_GPUS=2
+
+
+
+
+
+#cross vs reconstruction
+filelist_recon='dataset/filelists/voxceleb2_test_n_5000_reconstruction_5k.txt'
+filelist_cross='dataset/filelists/voxceleb2_test_n_5000_seed_797_cross_5K.txt' 
+if [ "$sample_mode" = "reconstruction" ]; then
+    sample_input_flags="--sampling_input_type=first_frame --sampling_ref_type=first_frame"
+    filelist=$filelist_recon
+elif [ "$sample_mode" = "cross" ]; then
+    sample_input_flags="--sampling_input_type=gt --sampling_ref_type=gt"
+    filelist=$filelist_cross
+else
+    echo "Error: sample_mode can only be \"cross\" or \"reconstruction\""
+    exit 0
+fi
+test_video_dir=$real_video_root
+mkdir -p $sample_path
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --learn_sigma True --num_channels 128 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm False"
+DIFFUSION_FLAGS="--predict_xstart False  --diffusion_steps 1000 --noise_schedule linear --rescale_timesteps False"
+SAMPLE_FLAGS="--sampling_seed=7   $sample_input_flags --timestep_respacing ddim25 --use_ddim True --model_path=$model_path --sample_path=$sample_path"
+DATA_FLAGS="--nframes 5 --nrefer 1 --image_size 128 --sampling_batch_size=32 "
+TFG_FLAGS="--face_hide_percentage 0.5 --use_ref=True --use_audio=True --audio_as_style=True"
+GEN_FLAGS="--generate_from_filelist 1 --test_video_dir=$test_video_dir --filelist=$filelist --save_orig=False --face_det_batch_size 64 --pads 0,0,0,0"
+
+if  [ "$NUM_GPUS" -gt 1 ]; then 
+    mpiexec -n $NUM_GPUS python generate_dist.py $MODEL_FLAGS  $DIFFUSION_FLAGS  $SAMPLE_FLAGS $DATA_FLAGS $TFG_FLAGS $GEN_FLAGS
+else
+    python generate.py $MODEL_FLAGS  $DIFFUSION_FLAGS  $SAMPLE_FLAGS $DATA_FLAGS $TFG_FLAGS $GEN_FLAGS
+fi

scripts/inference_single_video.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+#set paths and arguments
+sample_mode="cross" # or "reconstruction"
+NUM_GPUS=1
+generate_from_filelist=0
+video_path="path/to/video.mp4"
+audio_path="path/to/audio.mp4"
+out_path="path/to/output.mp4"
+model_path="path/to/model.pt"
+
+
+
+#cross vs reconstruction 
+if [ "$sample_mode" = "reconstruction" ]; then
+    sample_input_flags="--sampling_input_type=first_frame --sampling_ref_type=first_frame"
+elif [ "$sample_mode" = "cross" ]; then
+    sample_input_flags="--sampling_input_type=gt --sampling_ref_type=gt"
+else
+    echo "Error: sample_mode can only be \"cross\" or \"reconstruction\""
+    exit 0
+fi
+
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --learn_sigma True --num_channels 128 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm False"
+DIFFUSION_FLAGS="--predict_xstart False  --diffusion_steps 1000 --noise_schedule linear --rescale_timesteps False"
+SAMPLE_FLAGS="--sampling_seed=7   $sample_input_flags --timestep_respacing ddim25 --use_ddim True --model_path=$model_path"
+DATA_FLAGS="--nframes 5 --nrefer 1 --image_size 128 --sampling_batch_size=32 "
+TFG_FLAGS="--face_hide_percentage 0.5 --use_ref=True --use_audio=True --audio_as_style=True"
+GEN_FLAGS="--generate_from_filelist $generate_from_filelist  --video_path=$video_path --audio_path=$audio_path --out_path=$out_path --save_orig=False --face_det_batch_size 64 --pads 0,0,0,0 --is_voxceleb2=False"
+
+if  [ "$NUM_GPUS" -gt 1 ]; then 
+    mpiexec -n $NUM_GPUS python generate_dist.py $MODEL_FLAGS  $DIFFUSION_FLAGS  $SAMPLE_FLAGS $DATA_FLAGS $TFG_FLAGS $GEN_FLAGS
+else
+    python generate.py $MODEL_FLAGS  $DIFFUSION_FLAGS  $SAMPLE_FLAGS $DATA_FLAGS $TFG_FLAGS $GEN_FLAGS
+fi