Upload 3 files

README.md

@@ -1,11 +1,63 @@
----
-title: Diffused Heads
-emoji: 💻
-colorFrom: pink
-colorTo: green
-sdk: docker
-pinned: false
-license: cc-by-nc-sa-4.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Diffused Heads 
+
+Official repository for Diffused Heads: Diffusion Models Beat GANs on Talking-Face Generation.
+
+### [Project](https://mstypulkowski.github.io/diffusedheads/) | [Paper](https://arxiv.org/abs/2301.03396) | [Demo](https://youtu.be/DSipIDj-5q0)
+
+<p align="center">
+<img src='./intro.gif' width=400>
+</p>
+
+## Setup
+Python 3.x environment with [ffmpeg](https://www.ffmpeg.org/) is needed. The rest of the requirements can be installed using:
+```
+pip install -r requirements.txt
+```
+
+## Sampling
+Due to LRW license agreement, we are only able to provide a checkpoint of our model trained on CREMA.
+
+The entire test set generated by our method can be downloaded from [here](https://drive.google.com/file/d/1zWSqtV7O4WGkgh6WB55b8Mdg2lXXUudH/view?usp=drive_link).
+
+
+1. Download and unpack [checkpoints](https://drive.google.com/file/d/1U90egQvzERHclTYPCjZadrEMyF7TAPa-/view?usp=drive_link) (our model and pretrained audio encoder).
+   
+2. Download and unpack preprocessed CREMA [video](https://drive.google.com/file/d/1rM0FZLGiy-bJcxpv4CTlbUf0FuROubdk/view?usp=drive_link) and [audio](https://drive.google.com/file/d/1uS7Vi8EwarJFGQhsYHDMSkQmaNuiJIVW/view?usp=drive_link) files.
+   
+3. Specify paths and options in `config_crema.yaml` (check comments in the file).
+   
+4. Run the script
+```
+python sample.py
+```
+
+
+## Using your own data
+### Audio
+You can use audio recordings of your choosing freely. The only requirements are 16 kHz audio rate and a single audio channel. Please note our model is able to generate videos up to 9 seconds long depending on the audio.
+
+### Identity frame
+It is highly recommended to use a frame from the provided CREMA videos. This instance of the model was trained on clips with green background only. If you want to use your identity frame anyway, please follow this [repo](https://github.com/DinoMan/face-processor) for face alignment. Additionally, you may want to try segmenting the person and replacing background to green.
+
+## Training
+A training script will be uploaded in the future (ETA December 2023).
+
+## Citation
+```
+@article{stypulkowski2023diffused,
+  title={Diffused heads: Diffusion models beat gans on talking-face generation},
+  author={Stypu{\l}kowski, Micha{\l} and Vougioukas, Konstantinos and He, Sen and Zi{\k{e}}ba, Maciej and Petridis, Stavros and Pantic, Maja},
+  journal={arXiv preprint arXiv:2301.03396},
+  year={2023}
+}
+```
+
+## License
+This work is licensed under a
+[Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License][cc-by-nc-sa].
+
+[![CC BY-NC-SA 4.0][cc-by-nc-sa-image]][cc-by-nc-sa]
+
+[cc-by-nc-sa]: http://creativecommons.org/licenses/by-nc-sa/4.0/
+[cc-by-nc-sa-image]: https://licensebuttons.net/l/by-nc-sa/4.0/88x31.png
+[cc-by-nc-sa-shield]: https://img.shields.io/badge/License-CC%20BY--NC--SA%204.0-lightgrey.svg

diffusion.py

@@ -0,0 +1,145 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from tqdm import trange
+from torchvision.transforms import Compose
+
+
+class Diffusion(nn.Module):
+    def __init__(
+        self, nn_backbone, device, n_timesteps=1000, in_channels=3, image_size=128, out_channels=6, motion_transforms=None):
+        super(Diffusion, self).__init__()
+
+        self.nn_backbone = nn_backbone
+        self.n_timesteps = n_timesteps
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.x_shape = (image_size, image_size)
+        self.device = device
+
+        self.motion_transforms = motion_transforms if motion_transforms else Compose([])
+
+        self.timesteps = torch.arange(n_timesteps)
+        self.beta = self.get_beta_schedule()
+        self.set_params()
+        self.device = device
+
+    def sample(self, x_cond, audio_emb, n_audio_motion_embs=2, n_motion_frames=2, motion_channels=3):
+        with torch.no_grad():
+            n_frames = audio_emb.shape[1]
+
+            xT = torch.randn(x_cond.shape[0], n_frames, self.in_channels, self.x_shape[0], self.x_shape[1]).to(x_cond.device)
+
+            audio_ids = [0] * n_audio_motion_embs
+            for i in range(n_audio_motion_embs + 1):
+                audio_ids += [i]
+            
+            motion_frames = [self.motion_transforms(x_cond) for _ in range(n_motion_frames)]
+            motion_frames = torch.cat(motion_frames, dim=1)
+
+            samples = []
+            for i in trange(n_frames, desc=f'Sampling'):
+                sample_frame = self.sample_loop(xT[:, i].to(x_cond.device), x_cond, motion_frames, audio_emb[:, audio_ids])
+                samples.append(sample_frame.unsqueeze(1))
+                motion_frames = torch.cat([motion_frames[:, motion_channels:, :], self.motion_transforms(sample_frame)], dim=1)
+                audio_ids = audio_ids[1:] + [min(i + n_audio_motion_embs + 1, n_frames - 1)]
+        return torch.cat(samples, dim=1)
+
+    def sample_loop(self, xT, x_cond, motion_frames, audio_emb):
+        xt = xT
+        for i, t in reversed(list(enumerate(self.timesteps))):
+            timesteps = torch.tensor([t] * xT.shape[0]).to(xT.device)
+            timesteps_ids = torch.tensor([i] * xT.shape[0]).to(xT.device)
+            nn_out = self.nn_backbone(xt, timesteps, x_cond, motion_frames=motion_frames, audio_emb=audio_emb)
+            mean, logvar = self.get_p_params(xt, timesteps_ids, nn_out)
+            noise = torch.randn_like(xt) if t > 0 else torch.zeros_like(xt)
+            xt = mean + noise * torch.exp(logvar / 2)
+
+        return xt
+
+    def get_p_params(self, xt, timesteps, nn_out):
+        if self.in_channels == self.out_channels:
+            eps_pred = nn_out
+            p_logvar = self.expand(torch.log(self.beta[timesteps]))
+        else:
+            eps_pred, nu = nn_out.chunk(2, 1)
+            nu = (nu + 1) / 2
+            p_logvar = nu * self.expand(torch.log(self.beta[timesteps])) + (1 - nu) * self.expand(self.log_beta_tilde_clipped[timesteps])
+
+        p_mean, _ = self.get_q_params(xt, timesteps, eps_pred=eps_pred)
+        return p_mean, p_logvar
+
+    def get_q_params(self, xt, timesteps, eps_pred=None, x0=None):
+        if x0 is None:
+        # predict x0 from xt and eps_pred
+            coef1_x0 = self.expand(self.coef1_x0[timesteps])
+            coef2_x0 = self.expand(self.coef2_x0[timesteps])
+            x0 = coef1_x0 * xt - coef2_x0 * eps_pred
+            x0 = x0.clamp(-1, 1)
+
+        # q(x_{t-1} | x_t, x_0)
+        coef1_q = self.expand(self.coef1_q[timesteps])
+        coef2_q = self.expand(self.coef2_q[timesteps])
+        q_mean = coef1_q * x0 + coef2_q * xt
+
+        q_logvar = self.expand(self.log_beta_tilde_clipped[timesteps])
+
+        return q_mean, q_logvar
+
+    def get_beta_schedule(self, max_beta=0.999):
+        alpha_bar = lambda t: np.cos((t + 0.008) / 1.008 * np.pi / 2) ** 2
+        betas = []
+        for i in range(self.n_timesteps):
+            t1 = i / self.n_timesteps
+            t2 = (i + 1) / self.n_timesteps
+            betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+        return torch.tensor(betas).float()
+
+    def set_params(self):
+        self.alpha = 1 - self.beta
+        self.alpha_bar = torch.cumprod(self.alpha, dim=0)
+        self.alpha_bar_prev = torch.cat([torch.ones(1,), self.alpha_bar[:-1]])
+
+        self.beta_tilde = self.beta * (1.0 - self.alpha_bar_prev) / (1.0 - self.alpha_bar)
+        self.log_beta_tilde_clipped = torch.log(torch.cat([self.beta_tilde[1, None], self.beta_tilde[1:]]))
+
+        # to caluclate x0 from eps_pred
+        self.coef1_x0 = torch.sqrt(1.0 / self.alpha_bar)
+        self.coef2_x0 = torch.sqrt(1.0 / self.alpha_bar - 1)
+
+        # for q(x_{t-1} | x_t, x_0)
+        self.coef1_q = self.beta * torch.sqrt(self.alpha_bar_prev) / (1.0 - self.alpha_bar)
+        self.coef2_q = (1.0 - self.alpha_bar_prev) * torch.sqrt(self.alpha) / (1.0 - self.alpha_bar)
+
+    def space(self, n_timesteps_new):
+        # change parameters for spaced timesteps during sampling
+        self.timesteps = self.space_timesteps(self.n_timesteps, n_timesteps_new)
+        self.n_timesteps = n_timesteps_new
+
+        self.beta = self.get_spaced_beta()
+        self.set_params()
+
+    def space_timesteps(self, n_timesteps, target_timesteps):
+        all_steps = []
+        frac_stride = (n_timesteps - 1) / (target_timesteps - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(target_timesteps):
+            taken_steps.append(round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        return all_steps
+
+    def get_spaced_beta(self):
+        last_alpha_cumprod = 1.0
+        new_beta = []
+        for i, alpha_cumprod in enumerate(self.alpha_bar):
+            if i in self.timesteps:
+                new_beta.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+        return torch.tensor(new_beta)
+
+    def expand(self, arr, dim=4):
+        while arr.dim() < dim:
+            arr = arr[:, None]
+        return arr.to(self.device)

utils.py

@@ -0,0 +1,349 @@
+import os
+import tempfile
+import scipy.io.wavfile as wav
+import ffmpeg
+import cv2
+from PIL import Image
+
+import decord
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Compose, GaussianBlur, Grayscale, Resize
+import torchaudio
+
+decord.bridge.set_bridge('torch')
+torchaudio.set_audio_backend("sox_io")
+
+
+class AudioEncoder(nn.Module):
+    """
+    A PyTorch Module to encode audio data into a fixed-size vector 
+    (also known as an "embedding"). This can be useful for various machine 
+    learning tasks such as classification, similarity matching, etc.
+    """
+    def __init__(self, path):
+        """
+        Initialize the AudioEncoder object.
+
+        Args:
+            path (str): The file path where the pre-trained model is stored.
+        """
+        super().__init__()
+        self.model = torch.jit.load(path)
+        self.register_buffer('hidden', torch.zeros(2, 1, 256))
+
+    def forward(self, audio):
+        """
+        The forward method is where the actual encoding happens. Given an 
+        audio sample, this function returns its corresponding embedding.
+
+        Args:
+            audio (Tensor): A PyTorch tensor containing the audio data.
+        
+        Returns:
+            Tensor: The embedding of the given audio.
+        """
+        self.reset()
+        x = create_windowed_sequence(audio, 3200, cutting_stride=640, pad_samples=3200-640, cut_dim=1)
+        embs = []                                           
+        for i in range(x.shape[1]):
+            emb, _, self.hidden = self.model(x[:, i], torch.LongTensor([3200]), init_state=self.hidden)
+            embs.append(emb)
+        return torch.vstack(embs)
+
+    def reset(self):
+        """
+        Resets the hidden states in the model. Call this function 
+        before processing a new audio sample to ensure that there is 
+        no state carried over from the previous sample.
+        """
+        self.hidden = torch.zeros(2, 1, 256).to(self.hidden.device)
+
+
+def get_audio_emb(audio_path, checkpoint, device):
+    """
+    This function takes the path of an audio file, loads it into a 
+    PyTorch tensor, and returns its embedding.
+
+    Args:
+        audio_path (str): The file path of the audio to be loaded.
+        checkpoint (str): The file path of the pre-trained model.
+        device (str): The computing device ('cpu' or 'cuda').
+
+    Returns:
+        Tensor, Tensor: The original audio as a tensor and its corresponding embedding.
+    """
+    audio, audio_rate = torchaudio.load(audio_path, channels_first=False)
+    assert audio_rate == 16000, 'Only 16 kHZ audio is supported.'
+    audio = audio[None, None, :, 0].to(device)
+
+    audio_encoder = AudioEncoder(checkpoint).to(device)
+
+    emb = audio_encoder(audio)
+    return audio, emb
+
+
+def get_id_frame(path, random=False, resize=128):
+    """
+    Retrieves a frame from either a video or image file. This frame can 
+    serve as an identifier or reference for the video or image.
+
+    Args:
+        path (str): File path to the video or image.
+        random (bool): Whether to randomly select a frame from the video.
+        resize (int): The dimensions to which the frame should be resized.
+
+    Returns:
+        Tensor: The image frame as a tensor.
+    """
+    if path.endswith('.mp4'):
+        vr = decord.VideoReader(path)
+        if random:
+            idx = [np.random.randint(len(vr))]
+        else:
+            idx = [0]
+        frame = vr.get_batch(idx).permute(0, 3, 1, 2)
+    else:
+        frame = load_image_to_torch(path).unsqueeze(0)
+    
+    frame = (frame / 255) * 2 - 1
+    frame = Resize((resize, resize), antialias=True)(frame).float()
+    return frame
+
+
+def get_motion_transforms(args):
+    """
+    Applies a series of transformations like Gaussian blur and grayscale 
+    conversion based on the provided arguments. This is commonly used for 
+    data augmentation or preprocessing.
+
+    Args:
+        args (Namespace): Arguments containing options for motion transformations.
+    
+    Returns:
+        Compose: A composed function of transforms.
+    """
+    motion_transforms = []
+    if args.motion_blur:
+        motion_transforms.append(GaussianBlur(5, sigma=2.0))
+    if args.grayscale_motion:
+        motion_transforms.append(Grayscale(1))
+    return Compose(motion_transforms)
+
+
+def save_audio(path, audio, audio_rate=16000):
+    """
+    Saves the audio data as a WAV file.
+
+    Args:
+        path (str): The file path where the audio will be saved.
+        audio (Tensor or np.array): The audio data.
+        audio_rate (int): The sampling rate of the audio, defaults to 16000Hz.
+    """
+    if torch.is_tensor(audio):
+        aud = audio.squeeze().detach().cpu().numpy()
+    else:
+        aud = audio.copy()  # Make a copy so that we don't alter the object
+
+    aud = ((2 ** 15) * aud).astype(np.int16)
+    wav.write(path, audio_rate, aud)
+
+
+def save_video(path, video, fps=25, scale=2, audio=None, audio_rate=16000, overlay_pts=None, ffmpeg_experimental=False):
+    """
+    Saves the video data as an MP4 file. Optionally includes audio and overlay points.
+
+    Args:
+        path (str): The file path where the video will be saved.
+        video (Tensor or np.array): The video data.
+        fps (int): Frames per second of the video.
+        scale (int): Scaling factor for the video dimensions.
+        audio (Tensor or np.array, optional): Audio data.
+        audio_rate (int, optional): The sampling rate for the audio.
+        overlay_pts (list of points, optional): Points to overlay on the video frames.
+        ffmpeg_experimental (bool): Whether to use experimental ffmpeg options.
+
+    Returns:
+        bool: Success status.
+    """
+    if not os.path.exists(os.path.dirname(path)):
+        os.makedirs(os.path.dirname(path))
+    success = True    
+    out_size = (scale * video.shape[-1], scale * video.shape[-2])
+    video_path = get_temp_path(os.path.split(path)[0], ext=".mp4")
+    if torch.is_tensor(video):
+        vid = video.squeeze().detach().cpu().numpy()
+    else:
+        vid = video.copy()  # Make a copy so that we don't alter the object
+
+    if np.min(vid) < 0:
+        vid = 127 * vid + 127
+    elif np.max(vid) <= 1:
+        vid = 255 * vid
+
+    is_color = True
+    if vid.ndim == 3:
+        is_color = False
+
+    writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), float(fps), out_size, isColor=is_color)
+    for i, frame in enumerate(vid):
+        if is_color:
+            frame = cv2.cvtColor(np.rollaxis(frame, 0, 3), cv2.COLOR_RGB2BGR)
+
+        if scale != 1:
+            frame = cv2.resize(frame, out_size)
+
+        write_frame = frame.astype('uint8')
+
+        if overlay_pts is not None:
+            for pt in overlay_pts[i]:
+                cv2.circle(write_frame, (int(scale * pt[0]), int(scale * pt[1])), 2, (0, 0, 0), -1)
+
+        writer.write(write_frame)
+    writer.release()
+
+    inputs = [ffmpeg.input(video_path)['v']]
+
+    if audio is not None:  # Save the audio file
+        audio_path = swp_extension(video_path, ".wav")
+        save_audio(audio_path, audio, audio_rate)
+        inputs += [ffmpeg.input(audio_path)['a']]
+
+    try:
+        if ffmpeg_experimental:
+            out = ffmpeg.output(*inputs, path, strict='-2', loglevel="panic", vcodec='h264').overwrite_output()
+        else:
+            out = ffmpeg.output(*inputs, path, loglevel="panic", vcodec='h264').overwrite_output()
+        out.run(quiet=True)
+    except:
+        success = False
+
+    if audio is not None and os.path.isfile(audio_path):
+        os.remove(audio_path)
+    if os.path.isfile(video_path):
+        os.remove(video_path)
+
+    return success
+
+
+def load_image_to_torch(dir):
+    """
+    Load an image from disk and convert it to a PyTorch tensor.
+    
+    Args:
+        dir (str): The directory path to the image file.
+        
+    Returns:
+        torch.Tensor: A tensor representation of the image.
+    """
+    img = Image.open(dir).convert('RGB')
+    img = np.array(img)
+    return torch.from_numpy(img).permute(2, 0, 1)
+
+
+def get_temp_path(tmp_dir, mode="", ext=""):
+    """
+    Generate a temporary file path for storing data.
+    
+    Args:
+        tmp_dir (str): The directory where the temporary file will be created.
+        mode (str, optional): A string to append to the file name.
+        ext (str, optional): The file extension.
+        
+    Returns:
+        str: The full path to the temporary file.
+    """
+    file_path = next(tempfile._get_candidate_names()) + mode + ext
+    if not os.path.exists(tmp_dir):
+        os.makedirs(tmp_dir)
+    file_path = os.path.join(tmp_dir, file_path)
+    return file_path
+
+
+def swp_extension(file, ext):
+    """
+    Swap the extension of a given file name.
+    
+    Args:
+        file (str): The original file name.
+        ext (str): The new extension.
+        
+    Returns:
+        str: The file name with the new extension.
+    """
+    return os.path.splitext(file)[0] + ext
+
+
+def pad_both_ends(tensor, left, right, dim=0):
+    """
+    Pad a tensor on both ends along a specific dimension.
+    
+    Args:
+        tensor (torch.Tensor): The tensor to be padded.
+        left (int): The padding size for the left side.
+        right (int): The padding size for the right side.
+        dim (int, optional): The dimension along which to pad.
+        
+    Returns:
+        torch.Tensor: The padded tensor.
+    """
+    no_dims = len(tensor.size())
+    if dim == -1:
+        dim = no_dims - 1
+
+    padding = [0] * 2 * no_dims
+    padding[2 * (no_dims - dim - 1)] = left
+    padding[2 * (no_dims - dim - 1) + 1] = right
+    return F.pad(tensor, padding, "constant", 0)
+
+
+def cut_n_stack(seq, snip_length, cut_dim=0, cutting_stride=None, pad_samples=0):
+    """
+    Divide a sequence tensor into smaller snips and stack them.
+    
+    Args:
+        seq (torch.Tensor): The original sequence tensor.
+        snip_length (int): The length of each snip.
+        cut_dim (int, optional): The dimension along which to cut.
+        cutting_stride (int, optional): The stride length for cutting. Defaults to snip_length.
+        pad_samples (int, optional): Number of samples to pad at both ends.
+        
+    Returns:
+        torch.Tensor: A tensor containing the stacked snips.
+    """
+    if cutting_stride is None:
+        cutting_stride = snip_length
+
+    pad_left = pad_samples // 2
+    pad_right = pad_samples - pad_samples // 2
+
+    seq = pad_both_ends(seq, pad_left, pad_right, dim=cut_dim)
+
+    stacked = seq.narrow(cut_dim, 0, snip_length).unsqueeze(0)
+    iterations = (seq.size()[cut_dim] - snip_length) // cutting_stride + 1
+    for i in range(1, iterations):
+        stacked = torch.cat((stacked, seq.narrow(cut_dim, i * cutting_stride, snip_length).unsqueeze(0)))
+    return stacked
+
+
+def create_windowed_sequence(seqs, snip_length, cut_dim=0, cutting_stride=None, pad_samples=0):
+    """
+    Create a windowed sequence from a list of sequences.
+    
+    Args:
+        seqs (list of torch.Tensor): List of sequence tensors.
+        snip_length (int): The length of each snip.
+        cut_dim (int, optional): The dimension along which to cut.
+        cutting_stride (int, optional): The stride length for cutting. Defaults to snip_length.
+        pad_samples (int, optional): Number of samples to pad at both ends.
+        
+    Returns:
+        torch.Tensor: A tensor containing the windowed sequences.
+    """
+    windowed_seqs = []
+    for seq in seqs:
+        windowed_seqs.append(cut_n_stack(seq, snip_length, cut_dim, cutting_stride, pad_samples).unsqueeze(0))
+
+    return torch.cat(windowed_seqs)