first commit

WavLM.py

@@ -0,0 +1,854 @@
+# --------------------------------------------------------
+# WavLM: Large-Scale Self-Supervised  Pre-training  for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf)
+# Github source: https://github.com/microsoft/unilm/tree/master/wavlm
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+import math
+import logging
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import LayerNorm
+from einops import rearrange
+import requests
+from clint.textui import progress
+import os
+from WavLM_modules import (
+    Fp32GroupNorm,
+    Fp32LayerNorm,
+    GradMultiply,
+    MultiheadAttention,
+    SamePad,
+    init_bert_params,
+    get_activation_fn,
+    TransposeLast,
+    GLU_Linear,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class WavLM_wrapper(nn.Module):
+    def __init__(
+        self, model_size="Base+", feed_as_frames=True, merge_type="cat", model_path=None
+    ):
+        super().__init__()
+        assert model_size in ["Base+", "Large"]
+        if model_path is None:
+            model_path = os.path.join(
+                os.path.dirname(__file__), f"WavLM-{model_size}.pt"
+            )
+        if not os.path.exists(model_path):
+            self.download_model(model_path, model_size)
+        checkpoint = torch.load(model_path)
+        cfg = WavLMConfig(checkpoint["cfg"])
+        self.cfg = cfg
+        self.model = WavLM(cfg)
+        self.model.load_state_dict(checkpoint["model"])
+        self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+        self.code_size = 768 * 2 if merge_type == "cat" else 768
+        self.merge_type = merge_type
+        self.feed_as_frames = feed_as_frames
+
+    def download_model(self, out_path, size: str = "Base+"):
+        print("Downloading model...")
+        if size == "Base+":
+            url = "https://valle.blob.core.windows.net/share/wavlm/WavLM-Base+.pt?sv=2020-08-04&st=2023-03-01T07%3A51%3A05Z&se=2033-03-02T07%3A51%3A00Z&sr=c&sp=rl&sig=QJXmSJG9DbMKf48UDIU1MfzIro8HQOf3sqlNXiflY1I%3D"
+        else:
+            url = "https://valle.blob.core.windows.net/share/wavlm/WavLM-Large.pt?sv=2020-08-04&st=2023-03-01T07%3A51%3A05Z&se=2033-03-02T07%3A51%3A00Z&sr=c&sp=rl&sig=QJXmSJG9DbMKf48UDIU1MfzIro8HQOf3sqlNXiflY1I%3D"
+        r = requests.get(url, allow_redirects=True, stream=True)
+        with open(out_path, "wb") as f:
+            total_length = int(r.headers.get("content-length"))
+            for chunk in progress.bar(
+                r.iter_content(chunk_size=1024), expected_size=(total_length / 1024) + 1
+            ):
+                if chunk:
+                    f.write(chunk)
+                    f.flush()
+        print("Model downloaded to %s" % out_path)
+
+    def forward(self, x):
+        """
+        Args:
+            x: (batch, n_frames, audio_features)
+        """
+        T = x.shape[1]
+
+        if self.feed_as_frames:
+            x = rearrange(x, "b f d -> (b f) d")
+        else:
+            x = rearrange(x, "b ... -> b (...)")
+
+        if self.cfg.normalize:
+            x = torch.nn.functional.layer_norm(x, x.shape)
+
+        x = self.model.extract_features(x)[0]  # B, new_features, C
+        if self.feed_as_frames:
+            x = rearrange(x, "(b f) d c -> b f d c", f=T)
+        else:
+            x = torch.nn.functional.interpolate(
+                x.permute(0, 2, 1), T * 2, mode="nearest"
+            )
+            x = rearrange(x, "b c (f d) -> b f d c", d=2)
+
+        if self.merge_type == "cat":
+            if x.dim() == 3:
+                return rearrange(x, "b d c -> b (d c)")
+            return rearrange(x, "b f d c -> b f (d c)")
+        elif self.merge_type == "sum":
+            return x.sum(dim=-2)
+        elif self.merge_type == "mean":
+            return x.mean(dim=-2)
+        elif self.merge_type == "None":
+            return x
+        else:
+            raise NotImplementedError
+
+
+def compute_mask_indices(
+    shape: Tuple[int, int],
+    padding_mask: Optional[torch.Tensor],
+    mask_prob: float,
+    mask_length: int,
+    mask_type: str = "static",
+    mask_other: float = 0.0,
+    min_masks: int = 0,
+    no_overlap: bool = False,
+    min_space: int = 0,
+) -> np.ndarray:
+    """
+    Computes random mask spans for a given shape
+    Args:
+        shape: the the shape for which to compute masks.
+            should be of size 2 where first element is batch size and 2nd is timesteps
+        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
+        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
+            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
+            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
+        mask_type: how to compute mask lengths
+            static = fixed size
+            uniform = sample from uniform distribution [mask_other, mask_length*2]
+            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
+            poisson = sample from possion distribution with lambda = mask length
+        min_masks: minimum number of masked spans
+        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
+        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
+    """
+
+    bsz, all_sz = shape
+    mask = np.full((bsz, all_sz), False)
+
+    all_num_mask = int(
+        # add a random number for probabilistic rounding
+        mask_prob * all_sz / float(mask_length) + np.random.rand()
+    )
+
+    all_num_mask = max(min_masks, all_num_mask)
+
+    mask_idcs = []
+    for i in range(bsz):
+        if padding_mask is not None:
+            sz = all_sz - padding_mask[i].long().sum().item()
+            num_mask = int(
+                # add a random number for probabilistic rounding
+                mask_prob * sz / float(mask_length) + np.random.rand()
+            )
+            num_mask = max(min_masks, num_mask)
+        else:
+            sz = all_sz
+            num_mask = all_num_mask
+
+        if mask_type == "static":
+            lengths = np.full(num_mask, mask_length)
+        elif mask_type == "uniform":
+            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
+        elif mask_type == "normal":
+            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
+            lengths = [max(1, int(round(x))) for x in lengths]
+        elif mask_type == "poisson":
+            lengths = np.random.poisson(mask_length, size=num_mask)
+            lengths = [int(round(x)) for x in lengths]
+        else:
+            raise Exception("unknown mask selection " + mask_type)
+
+        if sum(lengths) == 0:
+            lengths[0] = min(mask_length, sz - 1)
+
+        if no_overlap:
+            mask_idc = []
+
+            def arrange(s, e, length, keep_length):
+                span_start = np.random.randint(s, e - length)
+                mask_idc.extend(span_start + i for i in range(length))
+
+                new_parts = []
+                if span_start - s - min_space >= keep_length:
+                    new_parts.append((s, span_start - min_space + 1))
+                if e - span_start - keep_length - min_space > keep_length:
+                    new_parts.append((span_start + length + min_space, e))
+                return new_parts
+
+            parts = [(0, sz)]
+            min_length = min(lengths)
+            for length in sorted(lengths, reverse=True):
+                lens = np.fromiter(
+                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
+                    np.int,
+                )
+                l_sum = np.sum(lens)
+                if l_sum == 0:
+                    break
+                probs = lens / np.sum(lens)
+                c = np.random.choice(len(parts), p=probs)
+                s, e = parts.pop(c)
+                parts.extend(arrange(s, e, length, min_length))
+            mask_idc = np.asarray(mask_idc)
+        else:
+            min_len = min(lengths)
+            if sz - min_len <= num_mask:
+                min_len = sz - num_mask - 1
+
+            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
+
+            mask_idc = np.asarray(
+                [
+                    mask_idc[j] + offset
+                    for j in range(len(mask_idc))
+                    for offset in range(lengths[j])
+                ]
+            )
+
+        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
+
+    min_len = min([len(m) for m in mask_idcs])
+    for i, mask_idc in enumerate(mask_idcs):
+        if len(mask_idc) > min_len:
+            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
+        mask[i, mask_idc] = True
+
+    return mask
+
+
+class WavLMConfig:
+    def __init__(self, cfg=None):
+        self.extractor_mode: str = "default"  # mode for feature extractor. default has a single group norm with d groups in the first conv block, whereas layer_norm has layer norms in every block (meant to use with normalize=True)
+        self.encoder_layers: int = 12  # num encoder layers in the transformer
+
+        self.encoder_embed_dim: int = 768  # encoder embedding dimension
+        self.encoder_ffn_embed_dim: int = 3072  # encoder embedding dimension for FFN
+        self.encoder_attention_heads: int = 12  # num encoder attention heads
+        self.activation_fn: str = "gelu"  # activation function to use
+
+        self.layer_norm_first: bool = False  # apply layernorm first in the transformer
+        self.conv_feature_layers: str = "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2"  # string describing convolutional feature extraction layers in form of a python list that contains [(dim, kernel_size, stride), ...]
+        self.conv_bias: bool = False  # include bias in conv encoder
+        self.feature_grad_mult: float = (
+            1.0  # multiply feature extractor var grads by this
+        )
+
+        self.normalize: bool = (
+            False  # normalize input to have 0 mean and unit variance during training
+        )
+
+        # dropouts
+        self.dropout: float = 0.1  # dropout probability for the transformer
+        self.attention_dropout: float = 0.1  # dropout probability for attention weights
+        self.activation_dropout: float = (
+            0.0  # dropout probability after activation in FFN
+        )
+        self.encoder_layerdrop: float = (
+            0.0  # probability of dropping a tarnsformer layer
+        )
+        self.dropout_input: float = (
+            0.0  # dropout to apply to the input (after feat extr)
+        )
+        self.dropout_features: float = (
+            0.0  # dropout to apply to the features (after feat extr)
+        )
+
+        # masking
+        self.mask_length: int = 10  # mask length
+        self.mask_prob: float = 0.65  # probability of replacing a token with mask
+        self.mask_selection: str = "static"  # how to choose mask length
+        self.mask_other: float = 0  # secondary mask argument (used for more complex distributions), see help in compute_mask_indicesh
+        self.no_mask_overlap: bool = False  # whether to allow masks to overlap
+        self.mask_min_space: int = (
+            1  # min space between spans (if no overlap is enabled)
+        )
+
+        # channel masking
+        self.mask_channel_length: int = 10  # length of the mask for features (channels)
+        self.mask_channel_prob: float = 0.0  # probability of replacing a feature with 0
+        self.mask_channel_selection: str = (
+            "static"  # how to choose mask length for channel masking
+        )
+        self.mask_channel_other: float = 0  # secondary mask argument (used for more complex distributions), see help in compute_mask_indices
+        self.no_mask_channel_overlap: bool = (
+            False  # whether to allow channel masks to overlap
+        )
+        self.mask_channel_min_space: int = (
+            1  # min space between spans (if no overlap is enabled)
+        )
+
+        # positional embeddings
+        self.conv_pos: int = (
+            128  # number of filters for convolutional positional embeddings
+        )
+        self.conv_pos_groups: int = (
+            16  # number of groups for convolutional positional embedding
+        )
+
+        # relative position embedding
+        self.relative_position_embedding: bool = (
+            False  # apply relative position embedding
+        )
+        self.num_buckets: int = 320  # number of buckets for relative position embedding
+        self.max_distance: int = (
+            1280  # maximum distance for relative position embedding
+        )
+        self.gru_rel_pos: bool = False  # apply gated relative position embedding
+
+        if cfg is not None:
+            self.update(cfg)
+
+    def update(self, cfg: dict):
+        self.__dict__.update(cfg)
+
+
+class WavLM(nn.Module):
+    def __init__(
+        self,
+        cfg: WavLMConfig,
+    ) -> None:
+        super().__init__()
+        logger.info(f"WavLM Config: {cfg.__dict__}")
+
+        self.cfg = cfg
+        feature_enc_layers = eval(cfg.conv_feature_layers)
+        self.embed = feature_enc_layers[-1][0]
+
+        self.feature_extractor = ConvFeatureExtractionModel(
+            conv_layers=feature_enc_layers,
+            dropout=0.0,
+            mode=cfg.extractor_mode,
+            conv_bias=cfg.conv_bias,
+        )
+
+        self.post_extract_proj = (
+            nn.Linear(self.embed, cfg.encoder_embed_dim)
+            if self.embed != cfg.encoder_embed_dim
+            else None
+        )
+
+        self.mask_prob = cfg.mask_prob
+        self.mask_selection = cfg.mask_selection
+        self.mask_other = cfg.mask_other
+        self.mask_length = cfg.mask_length
+        self.no_mask_overlap = cfg.no_mask_overlap
+        self.mask_min_space = cfg.mask_min_space
+
+        self.mask_channel_prob = cfg.mask_channel_prob
+        self.mask_channel_selection = cfg.mask_channel_selection
+        self.mask_channel_other = cfg.mask_channel_other
+        self.mask_channel_length = cfg.mask_channel_length
+        self.no_mask_channel_overlap = cfg.no_mask_channel_overlap
+        self.mask_channel_min_space = cfg.mask_channel_min_space
+
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+        self.dropout_features = nn.Dropout(cfg.dropout_features)
+
+        self.feature_grad_mult = cfg.feature_grad_mult
+
+        self.mask_emb = nn.Parameter(
+            torch.FloatTensor(cfg.encoder_embed_dim).uniform_()
+        )
+
+        self.encoder = TransformerEncoder(cfg)
+        self.layer_norm = LayerNorm(self.embed)
+
+    def apply_mask(self, x, padding_mask):
+        B, T, C = x.shape
+        if self.mask_prob > 0:
+            mask_indices = compute_mask_indices(
+                (B, T),
+                padding_mask,
+                self.mask_prob,
+                self.mask_length,
+                self.mask_selection,
+                self.mask_other,
+                min_masks=2,
+                no_overlap=self.no_mask_overlap,
+                min_space=self.mask_min_space,
+            )
+            mask_indices = torch.from_numpy(mask_indices).to(x.device)
+            x[mask_indices] = self.mask_emb
+        else:
+            mask_indices = None
+
+        if self.mask_channel_prob > 0:
+            mask_channel_indices = compute_mask_indices(
+                (B, C),
+                None,
+                self.mask_channel_prob,
+                self.mask_channel_length,
+                self.mask_channel_selection,
+                self.mask_channel_other,
+                no_overlap=self.no_mask_channel_overlap,
+                min_space=self.mask_channel_min_space,
+            )
+            mask_channel_indices = (
+                torch.from_numpy(mask_channel_indices)
+                .to(x.device)
+                .unsqueeze(1)
+                .expand(-1, T, -1)
+            )
+            x[mask_channel_indices] = 0
+
+        return x, mask_indices
+
+    def forward_padding_mask(
+        self,
+        features: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        extra = padding_mask.size(1) % features.size(1)
+        if extra > 0:
+            padding_mask = padding_mask[:, :-extra]
+        padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
+        padding_mask = padding_mask.all(-1)
+        return padding_mask
+
+    def extract_features(
+        self,
+        source: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        mask: bool = False,
+        ret_conv: bool = False,
+        output_layer: Optional[int] = None,
+        ret_layer_results: bool = False,
+    ):
+        if self.feature_grad_mult > 0:
+            features = self.feature_extractor(source)
+            if self.feature_grad_mult != 1.0:
+                features = GradMultiply.apply(features, self.feature_grad_mult)
+        else:
+            with torch.no_grad():
+                features = self.feature_extractor(source)
+
+        features = features.transpose(1, 2)
+        features = self.layer_norm(features)
+
+        if padding_mask is not None:
+            padding_mask = self.forward_padding_mask(features, padding_mask)
+
+        if self.post_extract_proj is not None:
+            features = self.post_extract_proj(features)
+
+        features = self.dropout_input(features)
+
+        if mask:
+            x, mask_indices = self.apply_mask(features, padding_mask)
+        else:
+            x = features
+
+        # feature: (B, T, D), float
+        # target: (B, T), long
+        # x: (B, T, D), float
+        # padding_mask: (B, T), bool
+        # mask_indices: (B, T), bool
+        x, layer_results = self.encoder(
+            x,
+            padding_mask=padding_mask,
+            layer=None if output_layer is None else output_layer - 1,
+        )
+
+        res = {
+            "x": x,
+            "padding_mask": padding_mask,
+            "features": features,
+            "layer_results": layer_results,
+        }
+
+        feature = res["features"] if ret_conv else res["x"]
+        if ret_layer_results:
+            feature = (feature, res["layer_results"])
+        return feature, res["padding_mask"]
+
+
+class ConvFeatureExtractionModel(nn.Module):
+    def __init__(
+        self,
+        conv_layers: List[Tuple[int, int, int]],
+        dropout: float = 0.0,
+        mode: str = "default",
+        conv_bias: bool = False,
+        conv_type: str = "default",
+    ):
+        super().__init__()
+
+        assert mode in {"default", "layer_norm"}
+
+        def block(
+            n_in,
+            n_out,
+            k,
+            stride,
+            is_layer_norm=False,
+            is_group_norm=False,
+            conv_bias=False,
+        ):
+            def make_conv():
+                conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
+                nn.init.kaiming_normal_(conv.weight)
+                return conv
+
+            assert not (is_layer_norm and is_group_norm), (
+                "layer norm and group norm are exclusive"
+            )
+
+            if is_layer_norm:
+                return nn.Sequential(
+                    make_conv(),
+                    nn.Dropout(p=dropout),
+                    nn.Sequential(
+                        TransposeLast(),
+                        Fp32LayerNorm(dim, elementwise_affine=True),
+                        TransposeLast(),
+                    ),
+                    nn.GELU(),
+                )
+            elif is_group_norm:
+                return nn.Sequential(
+                    make_conv(),
+                    nn.Dropout(p=dropout),
+                    Fp32GroupNorm(dim, dim, affine=True),
+                    nn.GELU(),
+                )
+            else:
+                return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())
+
+        self.conv_type = conv_type
+        if self.conv_type == "default":
+            in_d = 1
+            self.conv_layers = nn.ModuleList()
+            for i, cl in enumerate(conv_layers):
+                assert len(cl) == 3, "invalid conv definition: " + str(cl)
+                (dim, k, stride) = cl
+
+                self.conv_layers.append(
+                    block(
+                        in_d,
+                        dim,
+                        k,
+                        stride,
+                        is_layer_norm=mode == "layer_norm",
+                        is_group_norm=mode == "default" and i == 0,
+                        conv_bias=conv_bias,
+                    )
+                )
+                in_d = dim
+        elif self.conv_type == "conv2d":
+            in_d = 1
+            self.conv_layers = nn.ModuleList()
+            for i, cl in enumerate(conv_layers):
+                assert len(cl) == 3
+                (dim, k, stride) = cl
+
+                self.conv_layers.append(torch.nn.Conv2d(in_d, dim, k, stride))
+                self.conv_layers.append(torch.nn.ReLU())
+                in_d = dim
+        elif self.conv_type == "custom":
+            in_d = 1
+            idim = 80
+            self.conv_layers = nn.ModuleList()
+            for i, cl in enumerate(conv_layers):
+                assert len(cl) == 3
+                (dim, k, stride) = cl
+                self.conv_layers.append(
+                    torch.nn.Conv2d(in_d, dim, k, stride, padding=1)
+                )
+                self.conv_layers.append(torch.nn.LayerNorm([dim, idim]))
+                self.conv_layers.append(torch.nn.ReLU())
+                in_d = dim
+                if (i + 1) % 2 == 0:
+                    self.conv_layers.append(
+                        torch.nn.MaxPool2d(2, stride=2, ceil_mode=True)
+                    )
+                    idim = int(math.ceil(idim / 2))
+        else:
+            pass
+
+    def forward(self, x, mask=None):
+        # BxT -> BxCxT
+        x = x.unsqueeze(1)
+        if self.conv_type == "custom":
+            for conv in self.conv_layers:
+                if isinstance(conv, nn.LayerNorm):
+                    x = x.transpose(1, 2)
+                    x = conv(x).transpose(1, 2)
+                else:
+                    x = conv(x)
+            x = x.transpose(2, 3).contiguous()
+            x = x.view(x.size(0), -1, x.size(-1))
+        else:
+            for conv in self.conv_layers:
+                x = conv(x)
+            if self.conv_type == "conv2d":
+                b, c, t, f = x.size()
+                x = x.transpose(2, 3).contiguous().view(b, c * f, t)
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        self.dropout = args.dropout
+        self.embedding_dim = args.encoder_embed_dim
+
+        self.pos_conv = nn.Conv1d(
+            self.embedding_dim,
+            self.embedding_dim,
+            kernel_size=args.conv_pos,
+            padding=args.conv_pos // 2,
+            groups=args.conv_pos_groups,
+        )
+        dropout = 0
+        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
+        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
+        nn.init.constant_(self.pos_conv.bias, 0)
+
+        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
+        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())
+
+        if hasattr(args, "relative_position_embedding"):
+            self.relative_position_embedding = args.relative_position_embedding
+            self.num_buckets = args.num_buckets
+            self.max_distance = args.max_distance
+        else:
+            self.relative_position_embedding = False
+            self.num_buckets = 0
+            self.max_distance = 0
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerSentenceEncoderLayer(
+                    embedding_dim=self.embedding_dim,
+                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
+                    num_attention_heads=args.encoder_attention_heads,
+                    dropout=self.dropout,
+                    attention_dropout=args.attention_dropout,
+                    activation_dropout=args.activation_dropout,
+                    activation_fn=args.activation_fn,
+                    layer_norm_first=args.layer_norm_first,
+                    has_relative_attention_bias=(
+                        self.relative_position_embedding and i == 0
+                    ),
+                    num_buckets=self.num_buckets,
+                    max_distance=self.max_distance,
+                    gru_rel_pos=args.gru_rel_pos,
+                )
+                for i in range(args.encoder_layers)
+            ]
+        )
+
+        self.layer_norm_first = args.layer_norm_first
+        self.layer_norm = LayerNorm(self.embedding_dim)
+        self.layerdrop = args.encoder_layerdrop
+
+        self.apply(init_bert_params)
+
+    def forward(self, x, padding_mask=None, streaming_mask=None, layer=None):
+        x, layer_results = self.extract_features(x, padding_mask, streaming_mask, layer)
+
+        if self.layer_norm_first and layer is None:
+            x = self.layer_norm(x)
+
+        return x, layer_results
+
+    def extract_features(
+        self, x, padding_mask=None, streaming_mask=None, tgt_layer=None
+    ):
+        if padding_mask is not None:
+            x[padding_mask] = 0
+
+        x_conv = self.pos_conv(x.transpose(1, 2))
+        x_conv = x_conv.transpose(1, 2)
+        x += x_conv
+
+        if not self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        layer_results = []
+        z = None
+        if tgt_layer is not None:
+            layer_results.append((x, z))
+        r = None
+        pos_bias = None
+        for i, layer in enumerate(self.layers):
+            dropout_probability = np.random.random()
+            if not self.training or (dropout_probability > self.layerdrop):
+                x, z, pos_bias = layer(
+                    x,
+                    self_attn_padding_mask=padding_mask,
+                    need_weights=False,
+                    self_attn_mask=streaming_mask,
+                    pos_bias=pos_bias,
+                )
+            if tgt_layer is not None:
+                layer_results.append((x, z))
+            if i == tgt_layer:
+                r = x
+                break
+
+        if r is not None:
+            x = r
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+
+        return x, layer_results
+
+
+class TransformerSentenceEncoderLayer(nn.Module):
+    """
+    Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
+    models.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: float = 768,
+        ffn_embedding_dim: float = 3072,
+        num_attention_heads: float = 8,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        activation_dropout: float = 0.1,
+        activation_fn: str = "relu",
+        layer_norm_first: bool = False,
+        has_relative_attention_bias: bool = False,
+        num_buckets: int = 0,
+        max_distance: int = 0,
+        rescale_init: bool = False,
+        gru_rel_pos: bool = False,
+    ) -> None:
+        super().__init__()
+        # Initialize parameters
+        self.embedding_dim = embedding_dim
+        self.dropout = dropout
+        self.activation_dropout = activation_dropout
+
+        # Initialize blocks
+        self.activation_name = activation_fn
+        self.activation_fn = get_activation_fn(activation_fn)
+        self.self_attn = MultiheadAttention(
+            self.embedding_dim,
+            num_attention_heads,
+            dropout=attention_dropout,
+            self_attention=True,
+            has_relative_attention_bias=has_relative_attention_bias,
+            num_buckets=num_buckets,
+            max_distance=max_distance,
+            rescale_init=rescale_init,
+            gru_rel_pos=gru_rel_pos,
+        )
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(self.activation_dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.layer_norm_first = layer_norm_first
+
+        # layer norm associated with the self attention layer
+        self.self_attn_layer_norm = LayerNorm(self.embedding_dim)
+
+        if self.activation_name == "glu":
+            self.fc1 = GLU_Linear(self.embedding_dim, ffn_embedding_dim, "swish")
+        else:
+            self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
+
+        # layer norm associated with the position wise feed-forward NN
+        self.final_layer_norm = LayerNorm(self.embedding_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        self_attn_mask: torch.Tensor = None,
+        self_attn_padding_mask: torch.Tensor = None,
+        need_weights: bool = False,
+        pos_bias=None,
+    ):
+        """
+        LayerNorm is applied either before or after the self-attention/ffn
+        modules similar to the original Transformer imlementation.
+        """
+        residual = x
+
+        if self.layer_norm_first:
+            x = self.self_attn_layer_norm(x)
+            x, attn, pos_bias = self.self_attn(
+                query=x,
+                key=x,
+                value=x,
+                key_padding_mask=self_attn_padding_mask,
+                need_weights=False,
+                attn_mask=self_attn_mask,
+                position_bias=pos_bias,
+            )
+            x = self.dropout1(x)
+            x = residual + x
+
+            residual = x
+            x = self.final_layer_norm(x)
+            if self.activation_name == "glu":
+                x = self.fc1(x)
+            else:
+                x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual + x
+        else:
+            x, attn, pos_bias = self.self_attn(
+                query=x,
+                key=x,
+                value=x,
+                key_padding_mask=self_attn_padding_mask,
+                need_weights=need_weights,
+                attn_mask=self_attn_mask,
+                position_bias=pos_bias,
+            )
+
+            x = self.dropout1(x)
+            x = residual + x
+
+            x = self.self_attn_layer_norm(x)
+
+            residual = x
+            if self.activation_name == "glu":
+                x = self.fc1(x)
+            else:
+                x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual + x
+            x = self.final_layer_norm(x)
+
+        return x, attn, pos_bias

WavLM_modules.py

@@ -0,0 +1,765 @@
+# --------------------------------------------------------
+# WavLM: Large-Scale Self-Supervised  Pre-training  for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf)
+# Github source: https://github.com/microsoft/unilm/tree/master/wavlm
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+import math
+import warnings
+from typing import Dict, Optional, Tuple
+import torch
+from torch import Tensor, nn
+from torch.nn import Parameter
+import torch.nn.functional as F
+
+
+class TransposeLast(nn.Module):
+    def __init__(self, deconstruct_idx=None):
+        super().__init__()
+        self.deconstruct_idx = deconstruct_idx
+
+    def forward(self, x):
+        if self.deconstruct_idx is not None:
+            x = x[self.deconstruct_idx]
+        return x.transpose(-2, -1)
+
+
+class Fp32LayerNorm(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        output = F.layer_norm(
+            input.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(input)
+
+
+class Fp32GroupNorm(nn.GroupNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        output = F.group_norm(
+            input.float(),
+            self.num_groups,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(input)
+
+
+class GradMultiply(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, scale):
+        ctx.scale = scale
+        res = x.new(x)
+        return res
+
+    @staticmethod
+    def backward(ctx, grad):
+        return grad * ctx.scale, None
+
+
+class SamePad(nn.Module):
+    def __init__(self, kernel_size, causal=False):
+        super().__init__()
+        if causal:
+            self.remove = kernel_size - 1
+        else:
+            self.remove = 1 if kernel_size % 2 == 0 else 0
+
+    def forward(self, x):
+        if self.remove > 0:
+            x = x[:, :, : -self.remove]
+        return x
+
+
+class Swish(nn.Module):
+    """Swish function"""
+
+    def __init__(self):
+        """Construct an MultiHeadedAttention object."""
+        super(Swish, self).__init__()
+        self.act = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        return x * self.act(x)
+
+
+class GLU_Linear(nn.Module):
+    def __init__(self, input_dim, output_dim, glu_type="sigmoid", bias_in_glu=True):
+        super(GLU_Linear, self).__init__()
+
+        self.glu_type = glu_type
+        self.output_dim = output_dim
+
+        if glu_type == "sigmoid":
+            self.glu_act = torch.nn.Sigmoid()
+        elif glu_type == "swish":
+            self.glu_act = Swish()
+        elif glu_type == "relu":
+            self.glu_act = torch.nn.ReLU()
+        elif glu_type == "gelu":
+            self.glu_act = torch.nn.GELU()
+
+        if bias_in_glu:
+            self.linear = nn.Linear(input_dim, output_dim * 2, True)
+        else:
+            self.linear = nn.Linear(input_dim, output_dim * 2, False)
+
+    def forward(self, x):
+        # to be consistent with GLU_Linear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case
+        x = self.linear(x)
+
+        if self.glu_type == "bilinear":
+            x = x[:, :, 0 : self.output_dim] * x[:, :, self.output_dim : self.output_dim * 2]
+        else:
+            x = x[:, :, 0 : self.output_dim] * self.glu_act(x[:, :, self.output_dim : self.output_dim * 2])
+
+        return x
+
+
+def gelu_accurate(x):
+    if not hasattr(gelu_accurate, "_a"):
+        gelu_accurate._a = math.sqrt(2 / math.pi)
+    return 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3))))
+
+
+def gelu(x: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.gelu(x.float()).type_as(x)
+
+
+def get_activation_fn(activation: str):
+    """Returns the activation function corresponding to `activation`"""
+
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return gelu
+    elif activation == "gelu_fast":
+        warnings.warn("--activation-fn=gelu_fast has been renamed to gelu_accurate")
+        return gelu_accurate
+    elif activation == "gelu_accurate":
+        return gelu_accurate
+    elif activation == "tanh":
+        return torch.tanh
+    elif activation == "linear":
+        return lambda x: x
+    elif activation == "glu":
+        return lambda x: x
+    else:
+        raise RuntimeError("--activation-fn {} not supported".format(activation))
+
+
+def init_bert_params(module):
+    """
+    Initialize the weights specific to the BERT Model.
+    This overrides the default initializations depending on the specified arguments.
+        1. If normal_init_linear_weights is set then weights of linear
+           layer will be initialized using the normal distribution and
+           bais will be set to the specified value.
+        2. If normal_init_embed_weights is set then weights of embedding
+           layer will be initialized using the normal distribution.
+        3. If normal_init_proj_weights is set then weights of
+           in_project_weight for MultiHeadAttention initialized using
+           the normal distribution (to be validated).
+    """
+
+    def normal_(data):
+        # with FSDP, module params will be on CUDA, so we cast them back to CPU
+        # so that the RNG is consistent with and without FSDP
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        normal_(module.q_proj.weight.data)
+        normal_(module.k_proj.weight.data)
+        normal_(module.v_proj.weight.data)
+
+
+def quant_noise(module, p, block_size):
+    """
+    Wraps modules and applies quantization noise to the weights for
+    subsequent quantization with Iterative Product Quantization as
+    described in "Training with Quantization Noise for Extreme Model Compression"
+    Args:
+        - module: nn.Module
+        - p: amount of Quantization Noise
+        - block_size: size of the blocks for subsequent quantization with iPQ
+    Remarks:
+        - Module weights must have the right sizes wrt the block size
+        - Only Linear, Embedding and Conv2d modules are supported for the moment
+        - For more detail on how to quantize by blocks with convolutional weights,
+          see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
+        - We implement the simplest form of noise here as stated in the paper
+          which consists in randomly dropping blocks
+    """
+
+    # if no quantization noise, don't register hook
+    if p <= 0:
+        return module
+
+    # supported modules
+    assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
+
+    # test whether module.weight has the right sizes wrt block_size
+    is_conv = module.weight.ndim == 4
+
+    # 2D matrix
+    if not is_conv:
+        assert module.weight.size(1) % block_size == 0, "Input features must be a multiple of block sizes"
+
+    # 4D matrix
+    else:
+        # 1x1 convolutions
+        if module.kernel_size == (1, 1):
+            assert module.in_channels % block_size == 0, "Input channels must be a multiple of block sizes"
+        # regular convolutions
+        else:
+            k = module.kernel_size[0] * module.kernel_size[1]
+            assert k % block_size == 0, "Kernel size must be a multiple of block size"
+
+    def _forward_pre_hook(mod, input):
+        # no noise for evaluation
+        if mod.training:
+            if not is_conv:
+                # gather weight and sizes
+                weight = mod.weight
+                in_features = weight.size(1)
+                out_features = weight.size(0)
+
+                # split weight matrix into blocks and randomly drop selected blocks
+                mask = torch.zeros(in_features // block_size * out_features, device=weight.device)
+                mask.bernoulli_(p)
+                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
+
+            else:
+                # gather weight and sizes
+                weight = mod.weight
+                in_channels = mod.in_channels
+                out_channels = mod.out_channels
+
+                # split weight matrix into blocks and randomly drop selected blocks
+                if mod.kernel_size == (1, 1):
+                    mask = torch.zeros(
+                        int(in_channels // block_size * out_channels),
+                        device=weight.device,
+                    )
+                    mask.bernoulli_(p)
+                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
+                else:
+                    mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device)
+                    mask.bernoulli_(p)
+                    mask = mask.unsqueeze(2).unsqueeze(3).repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
+
+            # scale weights and apply mask
+            mask = mask.to(torch.bool)  # x.bool() is not currently supported in TorchScript
+            s = 1 / (1 - p)
+            mod.weight.data = s * weight.masked_fill(mask, 0)
+
+    module.register_forward_pre_hook(_forward_pre_hook)
+    return module
+
+
+class MultiheadAttention(nn.Module):
+    """Multi-headed attention.
+    See "Attention Is All You Need" for more details.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        self_attention=False,
+        encoder_decoder_attention=False,
+        q_noise=0.0,
+        qn_block_size=8,
+        has_relative_attention_bias=False,
+        num_buckets=32,
+        max_distance=128,
+        gru_rel_pos=False,
+        rescale_init=False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout_module = nn.Dropout(dropout)
+
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
+
+        self.head_dim = embed_dim // num_heads
+        self.q_head_dim = self.head_dim
+        self.k_head_dim = self.head_dim
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        k_bias = True
+        if rescale_init:
+            k_bias = False
+
+        k_embed_dim = embed_dim
+        q_embed_dim = embed_dim
+
+        self.k_proj = quant_noise(nn.Linear(self.kdim, k_embed_dim, bias=k_bias), q_noise, qn_block_size)
+        self.v_proj = quant_noise(nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size)
+        self.q_proj = quant_noise(nn.Linear(embed_dim, q_embed_dim, bias=bias), q_noise, qn_block_size)
+
+        self.out_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size)
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
+            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.gru_rel_pos = gru_rel_pos
+        if self.gru_rel_pos:
+            self.grep_linear = nn.Linear(self.q_head_dim, 8)
+            self.grep_a = nn.Parameter(torch.ones(1, num_heads, 1, 1))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.qkv_same_dim:
+            # Empirically observed the convergence to be much better with
+            # the scaled initialization
+            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+        else:
+            nn.init.xavier_uniform_(self.k_proj.weight)
+            nn.init.xavier_uniform_(self.v_proj.weight)
+            nn.init.xavier_uniform_(self.q_proj.weight)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            nn.init.xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            nn.init.xavier_normal_(self.bias_v)
+        if self.has_relative_attention_bias:
+            nn.init.xavier_normal_(self.relative_attention_bias.weight)
+
+    def _relative_positions_bucket(self, relative_positions, bidirectional=True):
+        num_buckets = self.num_buckets
+        max_distance = self.max_distance
+        relative_buckets = 0
+
+        if bidirectional:
+            num_buckets = num_buckets // 2
+            relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
+            relative_positions = torch.abs(relative_positions)
+        else:
+            relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions))
+
+        max_exact = num_buckets // 2
+        is_small = relative_positions < max_exact
+
+        relative_postion_if_large = max_exact + (
+            torch.log(relative_positions.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_postion_if_large = torch.min(
+            relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1)
+        )
+
+        relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large)
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length):
+        context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_positions_bucket(relative_position, bidirectional=True)
+        relative_position_bucket = relative_position_bucket.to(self.relative_attention_bias.weight.device)
+        values = self.relative_attention_bias(relative_position_bucket)
+        values = values.permute([2, 0, 1])
+        return values
+
+    def forward(
+        self,
+        query,
+        key: Optional[Tensor],
+        value: Optional[Tensor],
+        key_padding_mask: Optional[Tensor] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        need_weights: bool = True,
+        static_kv: bool = False,
+        attn_mask: Optional[Tensor] = None,
+        before_softmax: bool = False,
+        need_head_weights: bool = False,
+        position_bias: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+        """Input shape: Time x Batch x Channel
+        Args:
+            key_padding_mask (ByteTensor, optional): mask to exclude
+                keys that are pads, of shape `(batch, src_len)`, where
+                padding elements are indicated by 1s.
+            need_weights (bool, optional): return the attention weights,
+                averaged over heads (default: False).
+            attn_mask (ByteTensor, optional): typically used to
+                implement causal attention, where the mask prevents the
+                attention from looking forward in time (default: None).
+            before_softmax (bool, optional): return the raw attention
+                weights and values before the attention softmax.
+            need_head_weights (bool, optional): return the attention
+                weights for each head. Implies *need_weights*. Default:
+                return the average attention weights over all heads.
+        """
+        if need_head_weights:
+            need_weights = True
+
+        is_tpu = query.device.type == "xla"
+
+        tgt_len, bsz, embed_dim = query.size()
+        src_len = tgt_len
+        assert embed_dim == self.embed_dim
+        assert list(query.size()) == [tgt_len, bsz, embed_dim]
+        if key is not None:
+            src_len, key_bsz, _ = key.size()
+            if not torch.jit.is_scripting():
+                assert key_bsz == bsz
+                assert value is not None
+                assert src_len, bsz == value.shape[:2]
+
+        if self.has_relative_attention_bias and position_bias is None:
+            position_bias = self.compute_bias(tgt_len, src_len)
+            position_bias = (
+                position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1).view(bsz * self.num_heads, tgt_len, src_len)
+            )
+
+        if (
+            not is_tpu  # don't use PyTorch version on TPUs
+            and incremental_state is None
+            and not static_kv
+            # A workaround for quantization to work. Otherwise JIT compilation
+            # treats bias in linear module as method.
+            and not torch.jit.is_scripting()
+            and self.q_head_dim == self.head_dim
+        ):
+            assert key is not None and value is not None
+            assert attn_mask is None
+
+            attn_mask_rel_pos = None
+            if position_bias is not None:
+                attn_mask_rel_pos = position_bias
+                if self.gru_rel_pos:
+                    query_layer = query.transpose(0, 1)
+                    new_x_shape = query_layer.size()[:-1] + (self.num_heads, -1)
+                    query_layer = query_layer.view(*new_x_shape)
+                    query_layer = query_layer.permute(0, 2, 1, 3)
+                    _B, _H, _L, __ = query_layer.size()
+
+                    gate_a, gate_b = torch.sigmoid(
+                        self.grep_linear(query_layer).view(_B, _H, _L, 2, 4).sum(-1, keepdim=False)
+                    ).chunk(2, dim=-1)
+                    gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0
+                    attn_mask_rel_pos = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias
+
+                attn_mask_rel_pos = attn_mask_rel_pos.view((-1, tgt_len, tgt_len))
+            k_proj_bias = self.k_proj.bias
+            if k_proj_bias is None:
+                k_proj_bias = torch.zeros_like(self.q_proj.bias)
+
+            x, attn = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                torch.empty([0]),
+                torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout_module.p,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                self.training,
+                # self.training or self.dropout_module.apply_during_inference,
+                key_padding_mask,
+                need_weights,
+                attn_mask_rel_pos,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj.weight,
+                k_proj_weight=self.k_proj.weight,
+                v_proj_weight=self.v_proj.weight,
+            )
+            return x, attn, position_bias
+
+        if incremental_state is not None:
+            saved_state = self._get_input_buffer(incremental_state)
+            if saved_state is not None and "prev_key" in saved_state:
+                # previous time steps are cached - no need to recompute
+                # key and value if they are static
+                if static_kv:
+                    assert self.encoder_decoder_attention and not self.self_attention
+                    key = value = None
+        else:
+            saved_state = None
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        elif self.encoder_decoder_attention:
+            # encoder-decoder attention
+            q = self.q_proj(query)
+            if key is None:
+                assert value is None
+                k = v = None
+            else:
+                k = self.k_proj(key)
+                v = self.v_proj(key)
+
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+        q *= self.scaling
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
+                    ],
+                    dim=1,
+                )
+
+        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.q_head_dim).transpose(0, 1)
+        if k is not None:
+            k = k.contiguous().view(-1, bsz * self.num_heads, self.k_head_dim).transpose(0, 1)
+        if v is not None:
+            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+
+        if saved_state is not None:
+            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
+            if "prev_key" in saved_state:
+                _prev_key = saved_state["prev_key"]
+                assert _prev_key is not None
+                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    k = prev_key
+                else:
+                    assert k is not None
+                    k = torch.cat([prev_key, k], dim=1)
+                src_len = k.size(1)
+            if "prev_value" in saved_state:
+                _prev_value = saved_state["prev_value"]
+                assert _prev_value is not None
+                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    v = prev_value
+                else:
+                    assert v is not None
+                    v = torch.cat([prev_value, v], dim=1)
+            prev_key_padding_mask: Optional[Tensor] = None
+            if "prev_key_padding_mask" in saved_state:
+                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
+            assert k is not None and v is not None
+            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
+                key_padding_mask=key_padding_mask,
+                prev_key_padding_mask=prev_key_padding_mask,
+                batch_size=bsz,
+                src_len=k.size(1),
+                static_kv=static_kv,
+            )
+
+            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_key_padding_mask"] = key_padding_mask
+            # In this branch incremental_state is never None
+            assert incremental_state is not None
+            incremental_state = self._set_input_buffer(incremental_state, saved_state)
+        assert k is not None
+        assert k.size(1) == src_len
+
+        # This is part of a workaround to get around fork/join parallelism
+        # not supporting Optional types.
+        if key_padding_mask is not None and key_padding_mask.dim() == 0:
+            key_padding_mask = None
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.size(0) == bsz
+            assert key_padding_mask.size(1) == src_len
+
+        if self.add_zero_attn:
+            assert v is not None
+            src_len += 1
+            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
+            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
+            if attn_mask is not None:
+                attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask),
+                    ],
+                    dim=1,
+                )
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
+
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            # don't attend to padding symbols
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            if not is_tpu:
+                attn_weights = attn_weights.masked_fill(
+                    key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+                    float("-inf"),
+                )
+            else:
+                attn_weights = attn_weights.transpose(0, 2)
+                attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf"))
+                attn_weights = attn_weights.transpose(0, 2)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if before_softmax:
+            return attn_weights, v, position_bias
+
+        if position_bias is not None:
+            if self.gru_rel_pos == 1:
+                query_layer = q.view(bsz, self.num_heads, tgt_len, self.q_head_dim)
+                _B, _H, _L, __ = query_layer.size()
+                gate_a, gate_b = torch.sigmoid(
+                    self.grep_linear(query_layer).view(_B, _H, _L, 2, 4).sum(-1, keepdim=False)
+                ).chunk(2, dim=-1)
+                gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0
+                position_bias = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias
+
+            position_bias = position_bias.view(attn_weights.size())
+
+            attn_weights = attn_weights + position_bias
+
+        attn_weights_float = F.softmax(attn_weights, dim=-1)
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = self.dropout_module(attn_weights)
+
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+        attn_weights: Optional[Tensor] = None
+        if need_weights:
+            attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights = attn_weights.mean(dim=0)
+
+        return attn, attn_weights, position_bias
+
+    @staticmethod
+    def _append_prev_key_padding_mask(
+        key_padding_mask: Optional[Tensor],
+        prev_key_padding_mask: Optional[Tensor],
+        batch_size: int,
+        src_len: int,
+        static_kv: bool,
+    ) -> Optional[Tensor]:
+        # saved key padding masks have shape (bsz, seq_len)
+        if prev_key_padding_mask is not None and static_kv:
+            new_key_padding_mask = prev_key_padding_mask
+        elif prev_key_padding_mask is not None and key_padding_mask is not None:
+            new_key_padding_mask = torch.cat([prev_key_padding_mask.float(), key_padding_mask.float()], dim=1)
+        # During incremental decoding, as the padding token enters and
+        # leaves the frame, there will be a time when prev or current
+        # is None
+        elif prev_key_padding_mask is not None:
+            if src_len > prev_key_padding_mask.size(1):
+                filler = torch.zeros(
+                    (batch_size, src_len - prev_key_padding_mask.size(1)),
+                    device=prev_key_padding_mask.device,
+                )
+                new_key_padding_mask = torch.cat([prev_key_padding_mask.float(), filler.float()], dim=1)
+            else:
+                new_key_padding_mask = prev_key_padding_mask.float()
+        elif key_padding_mask is not None:
+            if src_len > key_padding_mask.size(1):
+                filler = torch.zeros(
+                    (batch_size, src_len - key_padding_mask.size(1)),
+                    device=key_padding_mask.device,
+                )
+                new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1)
+            else:
+                new_key_padding_mask = key_padding_mask.float()
+        else:
+            new_key_padding_mask = prev_key_padding_mask
+        return new_key_padding_mask
+
+    def _get_input_buffer(
+        self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
+    ) -> Dict[str, Optional[Tensor]]:
+        result = self.get_incremental_state(incremental_state, "attn_state")
+        if result is not None:
+            return result
+        else:
+            empty_result: Dict[str, Optional[Tensor]] = {}
+            return empty_result
+
+    def _set_input_buffer(
+        self,
+        incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
+        buffer: Dict[str, Optional[Tensor]],
+    ):
+        return self.set_incremental_state(incremental_state, "attn_state", buffer)
+
+    def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
+        return attn_weights

app.py

@@ -0,0 +1,978 @@
+import gradio as gr
+import torch
+import tempfile
+import os
+from vae_wrapper import VaeWrapper, encode_video_chunk
+from landmarks_extractor import LandmarksExtractor
+import decord
+from utils import (
+    get_raw_audio,
+    save_audio_video,
+    calculate_splits,
+    instantiate_from_config,
+    create_pipeline_inputs,
+)
+from transformers import HubertModel
+from einops import rearrange
+import numpy as np
+from WavLM import WavLM_wrapper
+from omegaconf import OmegaConf
+from inference_functions import (
+    sample_keyframes,
+    sample_interpolation,
+)
+from wordle_game import WordleGame
+import torch.cuda.amp as amp  # Import amp for mixed precision
+
+
+# Set default tensor type to float16 for faster computation
+if torch.cuda.is_available():
+    # torch.set_default_tensor_type(torch.cuda.FloatTensor)
+    # Enable TF32 precision for better performance on Ampere+ GPUs
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+# Cache for video and audio processing
+cache = {
+    "video": {
+        "path": None,
+        "embedding": None,
+        "frames": None,
+        "landmarks": None,
+    },
+    "audio": {
+        "path": None,
+        "raw_audio": None,
+        "hubert_embedding": None,
+        "wavlm_embedding": None,
+    },
+}
+
+# Create mixed precision scaler
+scaler = amp.GradScaler()
+
+
+def load_model(
+    config: str,
+    device: str = "cuda",
+    ckpt: str = None,
+):
+    """
+    Load a model from configuration.
+
+    Args:
+        config: Path to model configuration file
+        device: Device to load the model on
+        num_frames: Number of frames to process
+        input_key: Input key for the model
+        ckpt: Optional checkpoint path
+
+    Returns:
+        Tuple of (model, filter, batch size)
+    """
+    config = OmegaConf.load(config)
+
+    config["model"]["params"]["input_key"] = "latents"
+
+    if ckpt is not None:
+        config.model.params.ckpt_path = ckpt
+
+    with torch.device(device):
+        model = instantiate_from_config(config.model).to(device).eval()
+        # Convert model to half precision
+        if torch.cuda.is_available():
+            model = model.half()
+            model.first_stage_model = model.first_stage_model.float()
+            print("Converted model to FP16 precision")
+
+        # Compile model for faster inference
+        if torch.cuda.is_available():
+            try:
+                model = torch.compile(model)
+                print(f"Successfully compiled model with torch.compile()")
+            except Exception as e:
+                print(f"Warning: Failed to compile model: {e}")
+
+    return model
+
+
+# keyframe_model = KeyframeModel(device=device)
+# interpolation_model = InterpolationModel(device=device)
+vae_model = VaeWrapper("video")
+if torch.cuda.is_available():
+    vae_model = vae_model.half()  # Convert to half precision
+    try:
+        vae_model = torch.compile(vae_model)
+        print("Successfully compiled vae_model in FP16")
+    except Exception as e:
+        print(f"Warning: Failed to compile vae_model: {e}")
+
+hubert_model = HubertModel.from_pretrained("facebook/hubert-base-ls960").cuda()
+if torch.cuda.is_available():
+    hubert_model = hubert_model.half()  # Convert to half precision
+    try:
+        hubert_model = torch.compile(hubert_model)
+        print("Successfully compiled hubert_model in FP16")
+    except Exception as e:
+        print(f"Warning: Failed to compile hubert_model: {e}")
+
+wavlm_model = WavLM_wrapper(
+    model_size="Base+",
+    feed_as_frames=False,
+    merge_type="None",
+    model_path="/vol/paramonos2/projects/antoni/code/Personal/code_prep/keysync/pretrained_models/checkpoints/WavLM-Base+.pt",
+).cuda()
+if torch.cuda.is_available():
+    wavlm_model = wavlm_model.half()  # Convert to half precision
+    try:
+        wavlm_model = torch.compile(wavlm_model)
+        print("Successfully compiled wavlm_model in FP16")
+    except Exception as e:
+        print(f"Warning: Failed to compile wavlm_model: {e}")
+
+landmarks_extractor = LandmarksExtractor()
+keyframe_model = load_model(
+    config="/vol/paramonos2/projects/antoni/code/Personal/code_prep/keysync/scripts/sampling/configs/keyframe.yaml",
+    ckpt="/vol/paramonos2/projects/antoni/code/Personal/code_prep/keysync/pretrained_models/checkpoints/keyframe_dub.pt",
+)
+interpolation_model = load_model(
+    config="/vol/paramonos2/projects/antoni/code/Personal/code_prep/keysync/scripts/sampling/configs/interpolation.yaml",
+    ckpt="/vol/paramonos2/projects/antoni/code/Personal/code_prep/keysync/pretrained_models/checkpoints/interpolation_dub.pt",
+)
+keyframe_model.en_and_decode_n_samples_a_time = 2
+interpolation_model.en_and_decode_n_samples_a_time = 2
+
+# Default media paths
+DEFAULT_VIDEO_PATH = os.path.join(
+    os.path.dirname(__file__), "assets", "sample_video.mp4"
+)
+DEFAULT_AUDIO_PATH = os.path.join(
+    os.path.dirname(__file__), "assets", "sample_audio.wav"
+)
+
+
+@torch.no_grad()
+def compute_video_embedding(video_reader, min_len):
+    """Compute embeddings from video"""
+
+    total_frames = min_len
+
+    encoded = []
+    video_frames = []
+    chunk_size = 16
+    resolution = 512
+
+    # # Create a progress bar for Gradio
+    progress = gr.Progress()
+
+    # Calculate total chunks for progress tracking
+    total_chunks = (total_frames + chunk_size - 1) // chunk_size
+
+    for i, start_idx in enumerate(range(0, total_frames, chunk_size)):
+        # Update progress bar
+        progress(i / total_chunks, desc="Processing video chunks")
+
+        end_idx = min(start_idx + chunk_size, total_frames)
+        video_chunk = video_reader.get_batch(range(start_idx, end_idx))
+        # Interpolate video chunk to the target resolution
+        video_chunk = rearrange(video_chunk, "f h w c -> f c h w")
+        video_chunk = torch.nn.functional.interpolate(
+            video_chunk,
+            size=(resolution, resolution),
+            mode="bilinear",
+            align_corners=False,
+        )
+        video_chunk = rearrange(video_chunk, "f c h w -> f h w c")
+        video_frames.append(video_chunk)
+
+        # Convert chunk to FP16 if using CUDA
+        if torch.cuda.is_available():
+            video_chunk = video_chunk.half()
+
+        # Always use autocast for FP16 computation
+        with amp.autocast(enabled=True):
+            encoded.append(encode_video_chunk(vae_model, video_chunk, resolution))
+
+    encoded = torch.cat(encoded, dim=0)
+    video_frames = torch.cat(video_frames, dim=0)
+    video_frames = rearrange(video_frames, "f h w c -> f c h w")
+    torch.cuda.empty_cache()
+    return encoded, video_frames
+
+
+@torch.no_grad()
+def compute_hubert_embedding(raw_audio):
+    """Compute embeddings from audio"""
+    print(f"Computing audio embedding from {raw_audio.shape}")
+
+    audio = (
+        (raw_audio - raw_audio.mean()) / torch.sqrt(raw_audio.var() + 1e-7)
+    ).unsqueeze(0)
+    chunks = 16000 * 20
+
+    # Create a progress bar for Gradio
+    progress = gr.Progress()
+
+    # Get audio embeddings
+    audio_embeddings = []
+    splits = list(calculate_splits(audio, chunks))
+    total_splits = len(splits)
+
+    for i, chunk in enumerate(splits):
+        # Update progress bar
+        progress(i / total_splits, desc="Processing audio chunks")
+
+        # Convert audio chunk to half precision
+        if torch.cuda.is_available():
+            chunk_cuda = chunk.cuda().half()
+        else:
+            chunk_cuda = chunk.cuda()
+
+        # Always use autocast for FP16 computation
+        with amp.autocast(enabled=True):
+            hidden_states = hubert_model(chunk_cuda)[0]
+
+        audio_embeddings.append(hidden_states)
+    audio_embeddings = torch.cat(audio_embeddings, dim=1)
+
+    # audio_embeddings = self.model.wav2vec2(rearrange(audio_frames, "f s -> () (f s)"))[0]
+    if audio_embeddings.shape[1] % 2 != 0:
+        audio_embeddings = torch.cat(
+            [audio_embeddings, torch.zeros_like(audio_embeddings[:, :1])], dim=1
+        )
+    audio_embeddings = rearrange(audio_embeddings, "() (f d) c -> f d c", d=2)
+    torch.cuda.empty_cache()
+
+    return audio_embeddings
+
+
+@torch.no_grad()
+def compute_wavlm_embedding(raw_audio):
+    """Compute embeddings from audio"""
+    audio = rearrange(raw_audio, "(f s) -> f s", s=640)
+
+    if audio.shape[0] % 2 != 0:
+        audio = torch.cat([audio, torch.zeros(1, 640)], dim=0)
+    chunks = 500
+
+    # Create a progress bar for Gradio
+    progress = gr.Progress()
+
+    # Get audio embeddings
+    audio_embeddings = []
+    splits = list(calculate_splits(audio, chunks))
+    total_splits = len(splits)
+
+    for i, chunk in enumerate(splits):
+        # Update progress bar
+        progress(i / total_splits, desc="Processing audio chunks")
+
+        # Convert chunk to half precision
+        if torch.cuda.is_available():
+            chunk_cuda = chunk.unsqueeze(0).cuda().half()
+        else:
+            chunk_cuda = chunk.unsqueeze(0).cuda()
+
+        # Always use autocast for FP16 computation
+        with amp.autocast(enabled=True):
+            wavlm_hidden_states = wavlm_model(chunk_cuda).squeeze(0)
+
+        audio_embeddings.append(wavlm_hidden_states)
+    audio_embeddings = torch.cat(audio_embeddings, dim=0)
+
+    torch.cuda.empty_cache()
+
+    return audio_embeddings
+
+
+@torch.no_grad()
+def extract_video_landmarks(video_frames):
+    """Extract landmarks from video frames"""
+
+    # Create a progress bar for Gradio
+    progress = gr.Progress()
+
+    landmarks = []
+    batch_size = 10
+
+    for i in range(0, len(video_frames), batch_size):
+        # Update progress bar
+        progress(i / len(video_frames), desc="Extracting facial landmarks")
+
+        batch = video_frames[i : i + batch_size].cpu().float()
+        batch_landmarks = landmarks_extractor.extract_landmarks(batch)
+        landmarks.extend(batch_landmarks)
+
+    torch.cuda.empty_cache()
+
+    # Convert landmarks to a list of numpy arrays with consistent shape
+    processed_landmarks = []
+
+    expected_shape = (68, 2)  # Common shape for facial landmarks
+
+    # Process each landmark to ensure consistent shape
+    last_valid_landmark = None
+    for i, lm in enumerate(landmarks):
+        if lm is not None and isinstance(lm, np.ndarray) and lm.shape == expected_shape:
+            processed_landmarks.append(lm)
+            last_valid_landmark = lm
+        else:
+            # Print information about inconsistent landmarks
+            if lm is None:
+                print(f"Warning: Landmark at index {i} is None")
+            elif not isinstance(lm, np.ndarray):
+                print(
+                    f"Warning: Landmark at index {i} is not a numpy array, type: {type(lm)}"
+                )
+            elif lm.shape != expected_shape:
+                print(
+                    f"Warning: Landmark at index {i} has shape {lm.shape}, expected {expected_shape}"
+                )
+
+            # Replace invalid landmarks with the closest valid landmark if available
+            if last_valid_landmark is not None:
+                processed_landmarks.append(last_valid_landmark.copy())
+            else:
+                # If no valid landmark has been seen yet, look ahead for a valid one
+                found_future_valid = False
+                for future_lm in landmarks[i + 1 :]:
+                    if (
+                        future_lm is not None
+                        and isinstance(future_lm, np.ndarray)
+                        and future_lm.shape == expected_shape
+                    ):
+                        processed_landmarks.append(future_lm.copy())
+                        found_future_valid = True
+                        break
+
+                # If no valid landmark found in the future, use zeros
+                if not found_future_valid:
+                    processed_landmarks.append(np.zeros(expected_shape))
+
+    return np.array(processed_landmarks)
+
+
+@torch.no_grad()
+def sample(
+    audio_list,
+    gt_keyframes,
+    masks_keyframes,
+    to_remove,
+    test_keyframes_list,
+    num_frames,
+    device,
+    emb,
+    force_uc_zero_embeddings,
+    n_batch_keyframes,
+    n_batch,
+    test_interpolation_list,
+    audio_interpolation_list,
+    masks_interpolation,
+    gt_interpolation,
+    model_keyframes,
+    model,
+):
+    # Create a progress bar for Gradio
+    progress = gr.Progress()
+
+    condition = torch.zeros(1, 3, 512, 512).to(device)
+    if torch.cuda.is_available():
+        condition = condition.half()
+
+    audio_list = rearrange(audio_list, "(b t) c d  -> b t c d", t=num_frames)
+    gt_keyframes = rearrange(gt_keyframes, "(b t) c h w -> b t c h w", t=num_frames)
+    # Rearrange masks_keyframes and save locally
+    masks_keyframes = rearrange(
+        masks_keyframes, "(b t) c h w -> b t c h w", t=num_frames
+    )
+
+    # Convert to_remove into chunks of num_frames
+    to_remove_chunks = [
+        to_remove[i : i + num_frames] for i in range(0, len(to_remove), num_frames)
+    ]
+    test_keyframes_list = [
+        test_keyframes_list[i : i + num_frames]
+        for i in range(0, len(test_keyframes_list), num_frames)
+    ]
+
+    audio_cond = audio_list
+    if emb is not None:
+        embbedings = emb.unsqueeze(0).to(device)
+        if torch.cuda.is_available():
+            embbedings = embbedings.half()
+    else:
+        embbedings = None
+
+    # One batch of keframes is approximately 7 seconds
+    chunk_size = 2
+    complete_video = []
+    start_idx = 0
+    last_frame_z = None
+    last_frame_x = None
+    last_keyframe_idx = None
+    last_to_remove = None
+
+    total_chunks = (len(audio_cond) + chunk_size - 1) // chunk_size
+
+    for chunk_idx, chunk_start in enumerate(range(0, len(audio_cond), chunk_size)):
+        # Update progress bar
+        progress(chunk_idx / total_chunks, desc="Generating video")
+
+        # Clear GPU cache between chunks
+        torch.cuda.empty_cache()
+
+        chunk_end = min(chunk_start + chunk_size, len(audio_cond))
+
+        chunk_audio_cond = audio_cond[chunk_start:chunk_end].cuda()
+        if torch.cuda.is_available():
+            chunk_audio_cond = chunk_audio_cond.half()
+
+        chunk_gt_keyframes = gt_keyframes[chunk_start:chunk_end].cuda()
+        chunk_masks = masks_keyframes[chunk_start:chunk_end].cuda()
+
+        if torch.cuda.is_available():
+            chunk_gt_keyframes = chunk_gt_keyframes.half()
+            chunk_masks = chunk_masks.half()
+
+        test_keyframes_list_unwrapped = [
+            elem
+            for sublist in test_keyframes_list[chunk_start:chunk_end]
+            for elem in sublist
+        ]
+        to_remove_chunks_unwrapped = [
+            elem
+            for sublist in to_remove_chunks[chunk_start:chunk_end]
+            for elem in sublist
+        ]
+
+        if last_keyframe_idx is not None:
+            test_keyframes_list_unwrapped = [
+                last_keyframe_idx
+            ] + test_keyframes_list_unwrapped
+            to_remove_chunks_unwrapped = [last_to_remove] + to_remove_chunks_unwrapped
+
+        last_keyframe_idx = test_keyframes_list_unwrapped[-1]
+        last_to_remove = to_remove_chunks_unwrapped[-1]
+        # Find the first non-None keyframe in the chunk
+        first_keyframe = next(
+            (kf for kf in test_keyframes_list_unwrapped if kf is not None), None
+        )
+
+        # Find the last non-None keyframe in the chunk
+        last_keyframe = next(
+            (kf for kf in reversed(test_keyframes_list_unwrapped) if kf is not None),
+            None,
+        )
+
+        start_idx = next(
+            (
+                idx
+                for idx, comb in enumerate(test_interpolation_list)
+                if comb[0] == first_keyframe
+            ),
+            None,
+        )
+        end_idx = next(
+            (
+                idx
+                for idx, comb in enumerate(reversed(test_interpolation_list))
+                if comb[1] == last_keyframe
+            ),
+            None,
+        )
+
+        if start_idx is not None and end_idx is not None:
+            end_idx = (
+                len(test_interpolation_list) - 1 - end_idx
+            )  # Adjust for reversed enumeration
+        end_idx += 1
+        if start_idx is None:
+            break
+        if end_idx < start_idx:
+            end_idx = len(audio_interpolation_list)
+
+        audio_interpolation_list_chunk = audio_interpolation_list[start_idx:end_idx]
+        chunk_masks_interpolation = masks_interpolation[start_idx:end_idx]
+        gt_interpolation_chunks = gt_interpolation[start_idx:end_idx]
+
+        if torch.cuda.is_available():
+            audio_interpolation_list_chunk = [
+                chunk.half() for chunk in audio_interpolation_list_chunk
+            ]
+            chunk_masks_interpolation = [
+                chunk.half() for chunk in chunk_masks_interpolation
+            ]
+            gt_interpolation_chunks = [
+                chunk.half() for chunk in gt_interpolation_chunks
+            ]
+
+        progress(chunk_idx / total_chunks, desc="Generating keyframes")
+
+        # Always use autocast for FP16 computation
+        with amp.autocast(enabled=True):
+            samples_z = sample_keyframes(
+                model_keyframes,
+                chunk_audio_cond,
+                chunk_gt_keyframes,
+                chunk_masks,
+                condition.cuda(),
+                num_frames,
+                24,
+                0.0,
+                device,
+                embbedings.cuda() if embbedings is not None else None,
+                force_uc_zero_embeddings,
+                n_batch_keyframes,
+                0,
+                1.0,
+                None,
+                gt_as_cond=False,
+            )
+
+        if last_frame_x is not None:
+            # samples_x = torch.cat([last_frame_x.unsqueeze(0), samples_x], axis=0)
+            samples_z = torch.cat([last_frame_z.unsqueeze(0), samples_z], axis=0)
+
+        # last_frame_x = samples_x[-1]
+        last_frame_z = samples_z[-1]
+
+        progress(chunk_idx / total_chunks, desc="Interpolating frames")
+
+        # Always use autocast for FP16 computation
+        with amp.autocast(enabled=True):
+            vid = sample_interpolation(
+                model,
+                samples_z,
+                # samples_x,
+                audio_interpolation_list_chunk,
+                gt_interpolation_chunks,
+                chunk_masks_interpolation,
+                condition.cuda(),
+                num_frames,
+                device,
+                1,
+                24,
+                0.0,
+                force_uc_zero_embeddings,
+                n_batch,
+                chunk_size,
+                1.0,
+                None,
+                cut_audio=False,
+                to_remove=to_remove_chunks_unwrapped,
+            )
+
+        if chunk_start == 0:
+            complete_video = vid
+        else:
+            complete_video = np.concatenate([complete_video[:-1], vid], axis=0)
+
+    return complete_video
+
+
+def process_video(video_input, audio_input, max_num_seconds):
+    """Main processing function to generate synchronized video"""
+
+    # Display a message to the user about the processing time
+    gr.Info("Processing video. This may take a while...", duration=10)
+    gr.Info(
+        "If you're tired of waiting, try playing the Wordle game in the other tab!",
+        duration=10,
+    )
+
+    # Use default media if none provided
+    if video_input is None:
+        video_input = DEFAULT_VIDEO_PATH
+        print(f"Using default video: {DEFAULT_VIDEO_PATH}")
+
+    if audio_input is None:
+        audio_input = DEFAULT_AUDIO_PATH
+        print(f"Using default audio: {DEFAULT_AUDIO_PATH}")
+
+    try:
+        # Calculate hashes for cache keys
+        video_path_hash = video_input
+        audio_path_hash = audio_input
+
+        # Check if we need to recompute video embeddings
+        video_cache_hit = cache["video"]["path"] == video_path_hash
+        audio_cache_hit = cache["audio"]["path"] == audio_path_hash
+
+        if video_cache_hit and audio_cache_hit:
+            print("Using cached video and audio computations")
+            # Make copies of cached data to avoid modifying cache
+            video_embedding = cache["video"]["embedding"].clone()
+            video_frames = cache["video"]["frames"].clone()
+            video_landmarks = cache["video"]["landmarks"].copy()
+            raw_audio = cache["audio"]["raw_audio"].clone()
+            raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+            hubert_embedding = cache["audio"]["hubert_embedding"].clone()
+            wavlm_embedding = cache["audio"]["wavlm_embedding"].clone()
+
+            # Ensure all data is truncated to the same length if needed
+            min_len = min(
+                len(video_frames),
+                len(raw_audio),
+                len(hubert_embedding),
+                len(wavlm_embedding),
+            )
+            video_frames = video_frames[:min_len]
+            video_embedding = video_embedding[:min_len]
+            video_landmarks = video_landmarks[:min_len]
+            raw_audio = raw_audio[:min_len]
+            hubert_embedding = hubert_embedding[:min_len]
+            wavlm_embedding = wavlm_embedding[:min_len]
+            raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+
+        else:
+            # Process video if needed
+            if not video_cache_hit:
+                print("Computing video embeddings and landmarks")
+                video_reader = decord.VideoReader(video_input)
+                decord.bridge.set_bridge("torch")
+
+                if not audio_cache_hit:
+                    # Need to process audio to determine min_len
+                    raw_audio = get_raw_audio(audio_input, 16000)
+                    if len(raw_audio) == 0 or len(video_reader) == 0:
+                        raise ValueError("Empty audio or video input")
+
+                    min_len = min(len(raw_audio), len(video_reader))
+
+                    # Store full audio in cache
+                    cache["audio"]["path"] = audio_path_hash
+                    cache["audio"]["raw_audio"] = raw_audio.clone()
+
+                    # Create truncated copy for processing
+                    raw_audio = raw_audio[:min_len]
+                    raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+                else:
+                    # Use cached audio - make a copy
+                    if cache["audio"]["raw_audio"] is None:
+                        raise ValueError("Cached audio is None")
+
+                    raw_audio = cache["audio"]["raw_audio"].clone()
+                    if len(raw_audio) == 0 or len(video_reader) == 0:
+                        raise ValueError("Empty cached audio or video input")
+
+                    min_len = min(len(raw_audio), len(video_reader))
+
+                    # Create truncated copy for processing
+                    raw_audio = raw_audio[:min_len]
+                    raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+
+                # Compute video embeddings and landmarks - store full version in cache
+                video_embedding, video_frames = compute_video_embedding(
+                    video_reader, len(video_reader)
+                )
+                video_landmarks = extract_video_landmarks(video_frames)
+
+                # Update video cache with full versions
+                cache["video"]["path"] = video_path_hash
+                cache["video"]["embedding"] = video_embedding
+                cache["video"]["frames"] = video_frames
+                cache["video"]["landmarks"] = video_landmarks
+
+                # Create truncated copies for processing
+                video_embedding = video_embedding[:min_len]
+                video_frames = video_frames[:min_len]
+                video_landmarks = video_landmarks[:min_len]
+
+            else:
+                # Use cached video data - make copies
+                print("Using cached video computations")
+
+                if (
+                    cache["video"]["embedding"] is None
+                    or cache["video"]["frames"] is None
+                    or cache["video"]["landmarks"] is None
+                ):
+                    raise ValueError("One or more video cache entries are None")
+
+                if not audio_cache_hit:
+                    # New audio with cached video
+                    raw_audio = get_raw_audio(audio_input, 16000)
+                    if len(raw_audio) == 0:
+                        raise ValueError("Empty audio input")
+
+                    # Store full audio in cache
+                    cache["audio"]["path"] = audio_path_hash
+                    cache["audio"]["raw_audio"] = raw_audio.clone()
+
+                    # Make copies of video data
+                    video_embedding = cache["video"]["embedding"].clone()
+                    video_frames = cache["video"]["frames"].clone()
+                    video_landmarks = cache["video"]["landmarks"].copy()
+
+                    # Determine truncation length and create truncated copies
+                    min_len = min(len(raw_audio), len(video_frames))
+                    raw_audio = raw_audio[:min_len]
+                    raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+                    video_frames = video_frames[:min_len]
+                    video_embedding = video_embedding[:min_len]
+                    video_landmarks = video_landmarks[:min_len]
+                else:
+                    # Both video and audio are cached - should not reach here
+                    # as it's handled in the first if statement
+                    pass
+
+            # Process audio if needed
+            if not audio_cache_hit:
+                print("Computing audio embeddings")
+
+                # Compute audio embeddings with the truncated audio
+                hubert_embedding = compute_hubert_embedding(raw_audio_reshape)
+                wavlm_embedding = compute_wavlm_embedding(raw_audio_reshape)
+
+                # Update audio cache with full embeddings
+                # Note: raw_audio was already cached above
+                cache["audio"]["hubert_embedding"] = hubert_embedding.clone()
+                cache["audio"]["wavlm_embedding"] = wavlm_embedding.clone()
+            else:
+                # Use cached audio data - make copies
+                if (
+                    cache["audio"]["hubert_embedding"] is None
+                    or cache["audio"]["wavlm_embedding"] is None
+                ):
+                    raise ValueError(
+                        "One or more audio embedding cache entries are None"
+                    )
+
+                hubert_embedding = cache["audio"]["hubert_embedding"].clone()
+                wavlm_embedding = cache["audio"]["wavlm_embedding"].clone()
+
+                # Make sure embeddings match the truncated video length if needed
+                if "min_len" in locals() and (
+                    min_len < len(hubert_embedding) or min_len < len(wavlm_embedding)
+                ):
+                    hubert_embedding = hubert_embedding[:min_len]
+                    wavlm_embedding = wavlm_embedding[:min_len]
+
+        # Apply max_num_seconds limit if specified
+        if max_num_seconds > 0:
+            # Convert seconds to frames (assuming 25 fps)
+            max_frames = int(max_num_seconds * 25)
+
+            # Truncate all data to max_frames
+            video_embedding = video_embedding[:max_frames]
+            video_frames = video_frames[:max_frames]
+            video_landmarks = video_landmarks[:max_frames]
+            hubert_embedding = hubert_embedding[:max_frames]
+            wavlm_embedding = wavlm_embedding[:max_frames]
+            raw_audio = raw_audio[:max_frames]
+            raw_audio_reshape = rearrange(raw_audio, "f s -> (f s)")
+
+        # Validate shapes before proceeding
+        assert video_embedding.shape[0] == hubert_embedding.shape[0], (
+            f"Video embedding length ({video_embedding.shape[0]}) doesn't match Hubert embedding length ({hubert_embedding.shape[0]})"
+        )
+        assert video_embedding.shape[0] == wavlm_embedding.shape[0], (
+            f"Video embedding length ({video_embedding.shape[0]}) doesn't match WavLM embedding length ({wavlm_embedding.shape[0]})"
+        )
+        assert video_embedding.shape[0] == video_landmarks.shape[0], (
+            f"Video embedding length ({video_embedding.shape[0]}) doesn't match landmarks length ({video_landmarks.shape[0]})"
+        )
+
+        print(f"Hubert embedding shape: {hubert_embedding.shape}")
+        print(f"WavLM embedding shape: {wavlm_embedding.shape}")
+        print(f"Video embedding shape: {video_embedding.shape}")
+        print(f"Video landmarks shape: {video_landmarks.shape}")
+
+        # Create pipeline inputs for models
+        (
+            interpolation_chunks,
+            keyframe_chunks,
+            audio_interpolation_chunks,
+            audio_keyframe_chunks,
+            emb_cond,
+            masks_keyframe_chunks,
+            masks_interpolation_chunks,
+            to_remove,
+            audio_interpolation_idx,
+            audio_keyframe_idx,
+        ) = create_pipeline_inputs(
+            hubert_embedding,
+            wavlm_embedding,
+            14,
+            video_embedding,
+            video_landmarks,
+            overlap=1,
+            add_zero_flag=True,
+            mask_arms=None,
+            nose_index=28,
+        )
+
+        complete_video = sample(
+            audio_keyframe_chunks,
+            keyframe_chunks,
+            masks_keyframe_chunks,
+            to_remove,
+            audio_keyframe_idx,
+            14,
+            "cuda",
+            emb_cond,
+            [],
+            3,
+            3,
+            audio_interpolation_idx,
+            audio_interpolation_chunks,
+            masks_interpolation_chunks,
+            interpolation_chunks,
+            keyframe_model,
+            interpolation_model,
+        )
+
+        complete_audio = rearrange(
+            raw_audio[: complete_video.shape[0]], "f s -> () (f s)"
+        )
+
+        # 4. Convert frames to video and combine with audio
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
+            output_path = temp_video.name
+
+        print("Saving video to", output_path)
+
+        save_audio_video(complete_video, audio=complete_audio, save_path=output_path)
+        torch.cuda.empty_cache()
+        return output_path
+
+    except Exception as e:
+        raise e
+        print(f"Error processing video: {str(e)}")
+        return None
+
+
+def get_max_duration(video_input, audio_input):
+    """Get the maximum duration in seconds for the slider"""
+    try:
+        # Default to 60 seconds if files don't exist
+        if video_input is None or not os.path.exists(video_input):
+            video_input = DEFAULT_VIDEO_PATH
+
+        if audio_input is None or not os.path.exists(audio_input):
+            audio_input = DEFAULT_AUDIO_PATH
+
+        # Get video duration
+        video_reader = decord.VideoReader(video_input)
+        video_duration = len(video_reader) / video_reader.get_avg_fps()
+
+        # Get audio duration
+        raw_audio = get_raw_audio(audio_input, 16000)
+        audio_duration = len(raw_audio) / 25  # Assuming 25 fps
+
+        # Return the minimum of the two durations
+        return min(video_duration, audio_duration)
+    except Exception as e:
+        print(f"Error getting max duration: {str(e)}")
+        return 60  # Default to 60 seconds
+
+
+def new_game_click(state):
+    """Handle the 'New Game' button click."""
+    message = state.new_game()
+    feedback_history = state.get_feedback_history()
+    return state, feedback_history, message
+
+
+def submit_guess_click(guess, state):
+    """Handle the 'Submit Guess' button click."""
+    message = state.submit_guess(guess)
+    feedback_history = state.get_feedback_history()
+    return state, feedback_history, message
+
+
+# Create Gradio interface
+with gr.Blocks(title="Video Synchronization with Diffusion Models") as demo:
+    gr.Markdown("# Video Synchronization with Diffusion Models")
+    gr.Markdown(
+        "Upload a video and audio to create a synchronized video with the same visuals but synchronized to the new audio."
+    )
+
+    with gr.Tabs():
+        with gr.TabItem("Video Synchronization"):
+            with gr.Row():
+                with gr.Column():
+                    video_input = gr.Video(
+                        label="Input Video",
+                        value=DEFAULT_VIDEO_PATH
+                        if os.path.exists(DEFAULT_VIDEO_PATH)
+                        else None,
+                        width=512,
+                        height=512,
+                    )
+                    audio_input = gr.Audio(
+                        label="Input Audio",
+                        type="filepath",
+                        value=DEFAULT_AUDIO_PATH
+                        if os.path.exists(DEFAULT_AUDIO_PATH)
+                        else None,
+                    )
+
+                    max_duration = gr.State(value=60)  # Default max duration
+
+                    max_seconds_slider = gr.Slider(
+                        minimum=0,
+                        maximum=60,  # Will be updated dynamically
+                        value=0,
+                        step=1,
+                        label="Max Duration (seconds, 0 = full length)",
+                        info="Limit the processing duration (0 means use full length)",
+                    )
+
+                    process_button = gr.Button("Generate Synchronized Video")
+
+                with gr.Column("Output Video"):
+                    video_output = gr.Video(label="Output Video", width=512, height=512)
+
+            # Update slider max value when inputs change
+            def update_slider_max(video, audio):
+                max_dur = get_max_duration(video, audio)
+                return {"maximum": max_dur, "__type__": "update"}
+
+            video_input.change(
+                update_slider_max, [video_input, audio_input], [max_seconds_slider]
+            )
+            audio_input.change(
+                update_slider_max, [video_input, audio_input], [max_seconds_slider]
+            )
+
+            # Show Wordle message when processing starts and hide when complete
+            process_button.click(
+                fn=process_video,
+                inputs=[video_input, audio_input, max_seconds_slider],
+                outputs=video_output,
+            )
+
+        with gr.TabItem("Wordle Game"):
+            state = gr.State(WordleGame())  # Persist the WordleGame instance
+            guess_input = gr.Textbox(label="Your guess (5 letters)", max_length=5)
+            submit_btn = gr.Button("Submit Guess")
+            new_game_btn = gr.Button("New Game")
+            feedback_display = gr.HTML(label="Guesses")
+            message_display = gr.Textbox(
+                label="Message", interactive=False, value="Click 'New Game' to start."
+            )
+            # Connect the 'New Game' button
+            new_game_btn.click(
+                fn=new_game_click,
+                inputs=[state],
+                outputs=[state, feedback_display, message_display],
+            )
+            # Connect the 'Submit Guess' button
+            submit_btn.click(
+                fn=submit_guess_click,
+                inputs=[guess_input, state],
+                outputs=[state, feedback_display, message_display],
+            )
+
+    gr.Markdown("## How it works")
+    gr.Markdown("""
+    1. The system extracts embeddings and landmarks from the input video
+    2. Audio embeddings are computed from the input audio
+    3. A keyframe model generates key visual frames
+    4. An interpolation model creates a smooth video between keyframes
+    5. The final video is rendered with the new audio
+    """)
+
+if __name__ == "__main__":
+    demo.launch()

data_utils.py

@@ -0,0 +1,635 @@
+import torch
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+from functools import partial
+import math
+
+
+def get_size(img):
+    if isinstance(img, (np.ndarray, torch.Tensor)):
+        return img.shape[1::-1]
+    else:
+        return img.size
+
+
+def imresample(img, sz):
+    im_data = torch.nn.functional.interpolate(img, size=sz, mode="area")
+    return im_data
+
+
+def crop_resize(img, box, image_size):
+    if isinstance(img, np.ndarray):
+        img = img[box[1] : box[3], box[0] : box[2]]
+        out = cv2.resize(
+            img, (image_size, image_size), interpolation=cv2.INTER_AREA
+        ).copy()
+    elif isinstance(img, torch.Tensor):
+        img = img[box[1] : box[3], box[0] : box[2]]
+        out = (
+            imresample(
+                img.permute(2, 0, 1).unsqueeze(0).float(), (image_size, image_size)
+            )
+            .byte()
+            .squeeze(0)
+            .permute(1, 2, 0)
+        )
+    else:
+        out = img.crop(box).copy().resize((image_size, image_size), Image.BILINEAR)
+    return out
+
+
+def fixed_image_standardization(image_tensor):
+    processed_tensor = (image_tensor - 127.5) / 128.0
+    return processed_tensor
+
+
+def extract_face(img, landmarks, image_size=160, margin=0, postprocess=False):
+    """Extract face + margin from images given facial landmarks.
+
+    Arguments:
+        img {PIL.Image/torch.Tensor/np.ndarray} -- Input image(s) with shape (B, H, W, C)
+        landmarks {numpy.ndarray} -- Facial landmarks with shape (B, 68, 2)
+        image_size {int} -- Output image size in pixels. The image will be square.
+        margin {int} -- Margin to add to bounding box, in terms of pixels in the final image.
+        postprocess {bool} -- Whether to apply standardization
+
+    Returns:
+        torch.tensor -- tensor representing the extracted faces with shape (B, H, W, C)
+    """
+    # Calculate bounding boxes from landmarks for all faces in batch
+    x_min = np.min(landmarks, axis=1)[:, 0]  # Shape: (B,)
+    y_min = np.min(landmarks, axis=1)[:, 1]  # Shape: (B,)
+    x_max = np.max(landmarks, axis=1)[:, 0]  # Shape: (B,)
+    y_max = np.max(landmarks, axis=1)[:, 1]  # Shape: (B,)
+
+    # Calculate margin for top only
+    box_height = y_max - y_min
+    top_margin = margin * box_height / (image_size - margin)
+
+    # Create boxes for all faces
+    boxes = np.stack(
+        [
+            x_min,
+            np.maximum(y_min - top_margin, 0),  # Only add margin to top
+            x_max,
+            y_max,
+        ],
+        axis=1,
+    ).astype(int)  # Shape: (B, 4)
+
+    # Process each face in the batch
+    faces = []
+    for i in range(len(boxes)):
+        face = crop_resize(img[i], boxes[i], image_size)
+        faces.append(face)
+
+    faces = torch.stack(faces, dim=0)
+    faces = faces.float()
+
+    if postprocess:
+        faces = fixed_image_standardization(faces)
+
+    return faces
+
+
+def crop_mouth_region(images, landmarks, crop_size=96):
+    """
+    Takes a fixed-size square crop centered on the mouth region.
+
+    Parameters:
+    - images: tensor/array of shape (num_frames, height, width, channels) or (height, width, channels)
+    - landmarks: numpy array of shape (num_frames, 68, 2) or (68, 2)
+    - crop_size: size of the square crop (both height and width)
+    - padding: percentage of padding around the mouth region (0.0 to 1.0)
+
+    Returns:
+    - List of fixed-size crops or single crop if input is single image
+    """
+    # Handle single image case
+    single_image = False
+    if len(images.shape) == 3:
+        images = images[None]
+        landmarks = landmarks[None]
+        single_image = True
+
+    num_frames = len(images)
+    crops = []
+
+    # Mouth landmarks indices (48-67 for mouth region)
+    mouth_indices = range(48, 68)
+
+    for i in range(num_frames):
+        # Get mouth landmarks for current frame
+        mouth_landmarks = landmarks[i][mouth_indices]
+
+        # Find center of mouth
+        center_x = int(np.mean(mouth_landmarks[:, 0]))
+        center_y = int(np.mean(mouth_landmarks[:, 1]))
+
+        # Calculate crop boundaries
+        half_size = crop_size // 2
+        left = max(0, center_x - half_size)
+        right = min(images.shape[2], center_x + half_size)
+        top = max(0, center_y - half_size)
+        bottom = min(images.shape[1], center_y + half_size)
+
+        # Adjust if crop would go out of bounds
+        if left == 0:
+            right = crop_size
+        if right == images.shape[2]:
+            left = images.shape[2] - crop_size
+        if top == 0:
+            bottom = crop_size
+        if bottom == images.shape[1]:
+            top = images.shape[1] - crop_size
+
+        # Take the crop
+        crop = images[i, top:bottom, left:right]
+        crops.append(crop)
+
+    return crops[0] if single_image else crops
+
+
+def create_masks_from_landmarks_box(
+    landmark_list, img_shape, nose_index=28, dtype="uint8", box_expand=0.0
+):
+    height, width = img_shape[:2]
+    num_frames = landmark_list.shape[0]
+
+    # Initialize the masks array
+    masks = np.zeros((num_frames, height, width), dtype=dtype)
+
+    if 0 <= box_expand < 1:
+        box_expand = int(box_expand * width)
+
+    for i in range(num_frames):
+        # Get the landmarks for the current frame
+        landmarks = landmark_list[i]
+
+        # Get the y-coordinate of the nose landmark
+        nose_point_h = landmarks[nose_index, 1]
+        cut_h = nose_point_h
+
+        # Find the leftmost and rightmost landmarks
+        far_left_index = np.argmin(landmarks[:, 0])
+        far_right_index = np.argmax(landmarks[:, 0])
+
+        # Define the points for the mask contour
+        left_up_point = np.array(
+            [landmarks[far_left_index][0], cut_h - box_expand], dtype=np.int32
+        )
+        left_down_point = np.array(
+            [landmarks[far_left_index][0], height], dtype=np.int32
+        )
+        right_up_point = np.array(
+            [landmarks[far_right_index][0], cut_h - box_expand], dtype=np.int32
+        )
+        right_down_point = np.array(
+            [landmarks[far_right_index][0], height], dtype=np.int32
+        )
+
+        # Define the contour
+        contour = np.array(
+            [[left_up_point, left_down_point, right_down_point, right_up_point]]
+        )
+
+        # Draw the contour on the mask
+        cv2.drawContours(masks[i], [contour], -1, color=(1), thickness=cv2.FILLED)
+
+    return torch.from_numpy(masks)
+
+
+def create_masks_from_landmarks_full_size(
+    landmarks_batch,
+    image_height,
+    image_width,
+    start_index=48,
+    end_index=68,
+    offset=0,
+    nose_index=33,
+):
+    """
+    Efficiently creates a batch of masks using vectorized operations where each mask has ones from the highest
+    landmark in the specified range (adjusted by an offset) to the bottom of the image, and zeros otherwise.
+
+    Parameters:
+    - landmarks_batch (np.array): An array of shape (B, 68, 2) containing facial landmarks for multiple samples.
+    - image_height (int): The height of the image for which masks are created.
+    - image_width (int): The width of the image for which masks are created.
+    - start_index (int): The starting index of the range to check (inclusive).
+    - end_index (int): The ending index of the range to check (inclusive).
+    - offset (int): An offset to add or subtract from the y-coordinate of the highest landmark.
+
+    Returns:
+    - np.array: An array of masks of shape (B, image_height, image_width) for each batch.
+    """
+    # Extract the y-coordinates for the specified range across all batches
+    y_coords = landmarks_batch[:, nose_index : nose_index + 1, 1]
+
+    # Find the index of the minimum y-coordinate in the specified range for each batch
+    min_y_indices = np.argmin(y_coords, axis=1)
+
+    # Gather the highest landmarks' y-coordinates using the indices found
+    highest_y_coords = y_coords[np.arange(len(y_coords)), min_y_indices]
+
+    if abs(offset) < 1 and abs(offset) > 0:
+        offset = int(offset * image_height)
+
+    # Apply the offset to the highest y-coordinate
+    adjusted_y_coords = highest_y_coords + offset
+
+    # Clip the coordinates to stay within image boundaries
+    adjusted_y_coords = np.clip(adjusted_y_coords, 0, image_height - 1)
+
+    # Use broadcasting to create a mask without loops
+    # Create a range of indices from 0 to image_height - 1
+    all_indices = np.arange(image_height)
+
+    # Compare each index in 'all_indices' to each 'adjusted_y_coord' in the batch
+    # 'all_indices' has shape (image_height,), we reshape to (1, image_height) to broadcast against (B, 1)
+    mask_2d = (all_indices >= adjusted_y_coords[:, None]).astype(int)
+
+    # Extend the 2D mask to a full 3D mask of size (B, image_height, image_width)
+    full_mask = np.tile(mask_2d[:, :, np.newaxis], (1, 1, image_width))
+
+    return torch.from_numpy(full_mask)
+
+
+def expand_polygon(polygon, expand_size):
+    """
+    Expands the polygon outward by a specified number of pixels.
+
+    Parameters:
+    - polygon (list of tuples): The polygon points as (x, y).
+    - expand_size (int): The number of pixels to expand the polygon outward.
+
+    Returns:
+    - expanded_polygon (list of tuples): The expanded polygon points as (x, y).
+    """
+    if expand_size == 0:
+        return polygon
+
+    # Calculate centroid of the polygon
+    centroid_x = sum([point[0] for point in polygon]) / len(polygon)
+    centroid_y = sum([point[1] for point in polygon]) / len(polygon)
+
+    # Expand each point outward from the centroid
+    expanded_polygon = []
+    for x, y in polygon:
+        vector_x = x - centroid_x
+        vector_y = y - centroid_y
+        length = np.sqrt(vector_x**2 + vector_y**2)
+        if length == 0:
+            expanded_polygon.append((x, y))
+        else:
+            new_x = x + expand_size * (vector_x / length)
+            new_y = y + expand_size * (vector_y / length)
+            expanded_polygon.append((int(new_x), int(new_y)))
+
+    return expanded_polygon
+
+
+def create_masks_from_landmarks_mouth(
+    landmark_list, img_shape, nose_index=33, dtype="uint8", box_expand=0.0
+):
+    height, width = img_shape[:2]
+    num_frames = landmark_list.shape[0]
+
+    # Initialize the masks array
+    masks = np.zeros((num_frames, height, width), dtype=dtype)
+
+    if 0 <= box_expand < 1:
+        box_expand = int(box_expand * width)
+
+    for i in range(num_frames):
+        # Get the landmarks for the current frame
+        landmarks = landmark_list[i]
+
+        # Get the y-coordinate of the nose landmark
+        nose_point_h = landmarks[nose_index, 1]
+        cut_h = nose_point_h
+
+        # Find the leftmost and rightmost landmarks
+        far_left_index = np.argmin(landmarks[:, 0])
+        far_right_index = np.argmax(landmarks[:, 0])
+
+        # Find lowest landmark y-coordinate
+        lowest_y = np.max(landmarks[:, 1])
+        # Add box_expand to the lowest point
+        lowest_y = min(height, lowest_y + box_expand)
+
+        # Define the points for the mask contour
+        left_up_point = np.array(
+            [landmarks[far_left_index][0], cut_h - box_expand], dtype=np.int32
+        )
+        left_down_point = np.array(
+            [landmarks[far_left_index][0], lowest_y], dtype=np.int32
+        )
+        right_up_point = np.array(
+            [landmarks[far_right_index][0], cut_h - box_expand], dtype=np.int32
+        )
+        right_down_point = np.array(
+            [landmarks[far_right_index][0], lowest_y], dtype=np.int32
+        )
+
+        # Define the contour
+        contour = np.array(
+            [[left_up_point, left_down_point, right_down_point, right_up_point]]
+        )
+
+        # Draw the contour on the mask
+        cv2.drawContours(masks[i], [contour], -1, color=(1), thickness=cv2.FILLED)
+
+    return torch.from_numpy(masks)
+
+
+def create_face_mask_from_landmarks(
+    landmarks_batch, image_height, image_width, mask_expand=0
+):
+    """
+    Creates a batch of masks where each mask covers the face region using landmarks.
+
+    Parameters:
+    - landmarks_batch (np.array): An array of shape (B, 68, 2) containing facial landmarks for multiple samples.
+    - image_height (int): The height of the image for which masks are created.
+    - image_width (int): The width of the image for which masks are created.
+    - mask_expand (int): The number of pixels to expand the mask outward.
+
+    Returns:
+    - np.array: An array of masks of shape (B, image_height, image_width) for each batch.
+    """
+    # Initialize an array to hold all masks
+    masks = np.zeros(
+        (landmarks_batch.shape[0], image_height, image_width), dtype=np.uint8
+    )
+
+    if abs(mask_expand) < 1 and abs(mask_expand) > 0:
+        mask_expand = int(mask_expand * image_height)
+
+    for i, landmarks in enumerate(landmarks_batch):
+        # Create a blank image for each mask
+        mask = Image.new("L", (image_width, image_height), 0)
+        draw = ImageDraw.Draw(mask)
+
+        # Extract relevant landmarks for the face
+        jawline_landmarks = landmarks[2:15]  # Jawline
+        # upper_face_landmarks = landmarks[17:27]  # Eyebrows and top of nose bridge
+
+        # Combine landmarks to form a polygon around the face
+        # face_polygon = np.concatenate((jawline_landmarks, upper_face_landmarks[::-1]), axis=0)
+        face_polygon = jawline_landmarks
+
+        # Convert landmarks to a list of tuples
+        face_polygon = [(int(x), int(y)) for x, y in face_polygon]
+
+        # Expand the polygon if necessary
+        expanded_polygon = expand_polygon(face_polygon, mask_expand)
+
+        # Draw the polygon and fill it
+        draw.polygon(expanded_polygon, outline=1, fill=1)
+
+        # Convert mask to numpy array and add it to the batch of masks
+        masks[i] = np.array(mask)
+
+    return torch.from_numpy(masks)
+
+
+ALL_FIXED_POINTS = (
+    [i for i in range(0, 4)]
+    + [i for i in range(13, 17)]
+    + [i for i in range(27, 36)]
+    + [36, 39, 42, 45]
+)
+
+
+def gaussian_kernel(sigma, width, height):
+    """Create a 2D Gaussian kernel."""
+    x = torch.arange(0, width, 1) - width // 2
+    y = torch.arange(0, height, 1) - height // 2
+    x = x.float()
+    y = y.float()
+    x2 = x**2
+    y2 = y[:, None] ** 2
+    g = torch.exp(-(x2 + y2) / (2 * sigma**2))
+    return g / g.sum()
+
+
+def generate_hm(landmarks, height, width, n_points="all", sigma=3):
+    if n_points == "all":
+        Nlandmarks = range(len(landmarks))
+    elif n_points == "fixed":
+        Nlandmarks = ALL_FIXED_POINTS
+    elif n_points == "stable":
+        Nlandmarks = [33, 36, 39, 42, 45]
+
+    kernel = gaussian_kernel(sigma, width, height)
+    hm = torch.zeros((height, width))
+    for I in Nlandmarks:
+        x0, y0 = landmarks[I]
+        x0, y0 = int(x0), int(y0)
+        left, right = max(0, x0 - width // 2), min(width, x0 + width // 2)
+        top, bottom = max(0, y0 - height // 2), min(height, y0 + height // 2)
+        hm[top:bottom, left:right] += kernel[
+            max(0, -y0 + height // 2) : min(height, height - y0 + height // 2),
+            max(0, -x0 + width // 2) : min(width, width - x0 + width // 2),
+        ]
+    # Normalize the heatmap to have values between 0 and 1
+    max_val = hm.max()
+    if max_val > 0:
+        hm /= max_val
+    return hm
+
+
+def get_heatmap(landmarks, image_size, or_im_size, n_points="stable", sigma=4):
+    stack = []
+    seq_length = landmarks.shape[0]
+    if or_im_size[0] != image_size[0] or or_im_size[1] != image_size[1]:
+        landmarks = scale_landmarks(landmarks, or_im_size, image_size)
+    gen_single_heatmap = partial(
+        generate_hm,
+        height=image_size[0],
+        width=image_size[1],
+        n_points=n_points,
+        sigma=sigma,
+    )
+    for i in range(seq_length):
+        stack.append(gen_single_heatmap(landmarks[i]))
+
+    return torch.stack(stack, axis=0).unsqueeze(0)  # (1, seq_length, height, width)
+
+
+def scale_landmarks(landmarks, original_size, target_size):
+    """
+    Scale landmarks from original size to target size.
+
+    Parameters:
+    - landmarks (np.array): An array of shape (N, 2) containing facial landmarks.
+    - original_size (tuple): The size (height, width) for which the landmarks are currently scaled.
+    - target_size (tuple): The size (height, width) to which landmarks should be scaled.
+
+    Returns:
+    - scaled_landmarks (np.array): Scaled landmarks.
+    """
+    scale_y = target_size[0] / original_size[0]
+    scale_x = target_size[1] / original_size[1]
+    scaled_landmarks = landmarks * np.array([scale_x, scale_y])
+    return scaled_landmarks.astype(int)
+
+
+def draw_kps_image(
+    image_shape,
+    original_size,
+    landmarks,
+    color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255)],
+    rgb=True,
+    pts_width=4,
+):
+    stick_width = pts_width
+    limb_seq = np.array([[0, 2], [1, 2]])
+    kps = landmarks[[36, 45, 33], :]
+    kps = scale_landmarks(kps, original_size, image_shape)
+    if not rgb:  # Grayscale image
+        canvas = np.zeros((image_shape[0], image_shape[1], 1))
+        color_mode = "grayscale"
+    else:  # Color image
+        canvas = np.zeros((image_shape[0], image_shape[1], 3))
+        color_mode = "color"
+
+    polygon_cache = {}
+
+    for index in limb_seq:
+        color = color_list[index[0]]
+        if color_mode == "grayscale":
+            color = (
+                int(0.299 * color[2] + 0.587 * color[1] + 0.114 * color[0]),
+            )  # Convert to grayscale intensity
+
+        x = kps[index][:, 0]
+        y = kps[index][:, 1]
+        length = np.sqrt((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2)
+        angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1]))
+
+        cache_key = (
+            color,
+            int(np.mean(x)),
+            int(np.mean(y)),
+            int(length / 2),
+            int(angle),
+        )
+        if cache_key not in polygon_cache:
+            polygon_cache[cache_key] = cv2.ellipse2Poly(
+                (int(np.mean(x)), int(np.mean(y))),
+                (int(length / 2), stick_width),
+                int(angle),
+                0,
+                360,
+                1,
+            )
+
+        polygon = polygon_cache[cache_key]
+        cv2.fillConvexPoly(canvas, polygon, [int(c * 0.6) for c in color])
+
+    for idx, kp in enumerate(kps):
+        if color_mode == "grayscale":
+            color = (
+                int(
+                    0.299 * color_list[idx][2]
+                    + 0.587 * color_list[idx][1]
+                    + 0.114 * color_list[idx][0]
+                ),
+            )
+        else:
+            color = color_list[idx]
+        cv2.circle(canvas, (int(kp[0]), int(kp[1])), pts_width, color, -1)
+
+    return canvas.transpose(2, 0, 1)
+
+
+def create_landmarks_image(
+    landmarks,
+    original_size=(772, 772),
+    target_size=(772, 772),
+    point_size=3,
+    n_points="all",
+    dim=3,
+):
+    """
+    Creates an image of landmarks on a black background using efficient NumPy operations.
+
+    Parameters:
+    - landmarks (np.array): An array of shape (68, 2) containing facial landmarks.
+    - image_size (tuple): The size of the output image (height, width).
+    - point_size (int): The radius of each landmark point in pixels.
+
+    Returns:
+    - img (np.array): An image array with landmarks plotted.
+    """
+    if n_points == "all":
+        indexes = range(len(landmarks))
+    elif n_points == "fixed":
+        indexes = ALL_FIXED_POINTS
+    elif n_points == "stable":
+        indexes = [33, 36, 39, 42, 45]
+
+    landmarks = landmarks[indexes]
+
+    img = np.zeros(target_size, dtype=np.uint8)
+
+    landmarks = scale_landmarks(landmarks, original_size, target_size)
+
+    # Ensure the landmarks are in bounds and integer
+    landmarks = np.clip(
+        landmarks, [0, 0], [target_size[1] - 1, target_size[0] - 1]
+    ).astype(int)
+
+    # Get x and y coordinates from landmarks
+    x, y = landmarks[:, 0], landmarks[:, 1]
+
+    # Define a grid offset based on point_size around each landmark
+    offset = np.arange(-point_size // 2, point_size // 2 + 1)
+    grid_x, grid_y = np.meshgrid(offset, offset, indexing="ij")
+
+    # Calculate the full set of x and y coordinates for the points
+    full_x = x[:, np.newaxis, np.newaxis] + grid_x[np.newaxis, :, :]
+    full_y = y[:, np.newaxis, np.newaxis] + grid_y[np.newaxis, :, :]
+
+    # Clip the coordinates to stay within image boundaries
+    full_x = np.clip(full_x, 0, target_size[1] - 1)
+    full_y = np.clip(full_y, 0, target_size[0] - 1)
+
+    # Flatten the arrays to use them as indices
+    full_x = full_x.ravel()
+    full_y = full_y.ravel()
+
+    # Set the points in the image
+    img[full_y, full_x] = 255
+
+    return np.stack([img] * dim, axis=0)
+
+
+def trim_pad_audio(audio, sr, max_len_sec=None, max_len_raw=None):
+    len_file = audio.shape[-1]
+
+    if max_len_sec or max_len_raw:
+        max_len = max_len_raw if max_len_raw is not None else int(max_len_sec * sr)
+        if len_file < int(max_len):
+            # dummy = np.zeros((1, int(max_len_sec * sr) - len_file))
+            # extened_wav = np.concatenate((audio_data, dummy[0]))
+            extened_wav = torch.nn.functional.pad(
+                audio, (0, int(max_len) - len_file), "constant"
+            )
+        else:
+            extened_wav = audio[:, : int(max_len)]
+    else:
+        extened_wav = audio
+
+    return extened_wav
+
+
+def ssim_to_bin(ssim_score):
+    # Normalize the SSIM score to a 0-100 scale
+    normalized_diff_ssim = (1 - ((ssim_score + 1) / 2)) * 100
+    # Assign to one of the 100 bins
+    bin_index = float(min(np.floor(normalized_diff_ssim), 99))
+    return bin_index

inference_functions.py

@@ -0,0 +1,493 @@
+import torch
+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+from einops import rearrange, repeat
+import math
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    return list(set([x.input_key for x in conditioner.embedders]))
+
+
+def get_batch(keys, value_dict, N, T, device):
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "fps_id":
+            batch[key] = (
+                torch.tensor([value_dict["fps_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "motion_bucket_id":
+            batch[key] = (
+                torch.tensor([value_dict["motion_bucket_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "cond_aug":
+            batch[key] = repeat(
+                torch.tensor([value_dict["cond_aug"]]).to(device),
+                "1 -> b",
+                b=math.prod(N),
+            )
+        elif key == "cond_frames":
+            batch[key] = repeat(value_dict["cond_frames"], "1 ... -> b ...", b=N[0])
+        elif key == "cond_frames_without_noise":
+            batch[key] = repeat(
+                value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]
+            )
+        else:
+            batch[key] = value_dict[key]
+
+    if T is not None:
+        batch["num_video_frames"] = T
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+    return batch, batch_uc
+
+
+def merge_overlapping_segments(segments: torch.Tensor, overlap: int) -> torch.Tensor:
+    """
+    Merges overlapping segments by averaging overlapping frames.
+    Segments have shape (b, t, ...), where 'b' is the number of segments,
+    't' is frames per segment, and '...' are other dimensions.
+
+    Args:
+        segments: Tensor of shape (b, t, ...)
+        overlap: Integer, number of frames that overlap between consecutive segments
+
+    Returns:
+        Tensor of the merged video
+    """
+    # Get the shape details
+    b, t, *other_dims = segments.shape
+    num_frames = (b - 1) * (
+        t - overlap
+    ) + t  # Calculate the total number of frames in the merged video
+
+    # Initialize the output tensor and a count tensor to keep track of contributions for averaging
+    output_shape = [num_frames] + other_dims
+    output = torch.zeros(output_shape, dtype=segments.dtype, device=segments.device)
+    count = torch.zeros(output_shape, dtype=torch.float32, device=segments.device)
+
+    current_index = 0
+    for i in range(b):
+        end_index = current_index + t
+        # Add the segment to the output tensor
+        output[current_index:end_index] += rearrange(segments[i], "... -> ...")
+        # Increment the count tensor for each frame that's added
+        count[current_index:end_index] += 1
+        # Update the starting index for the next segment
+        current_index += t - overlap
+
+    # Avoid division by zero
+    count[count == 0] = 1
+    # Average the frames where there's overlap
+    output /= count
+
+    return output
+
+
+def get_batch_overlap(
+    keys: List[str],
+    value_dict: Dict[str, Any],
+    N: Tuple[int, ...],
+    T: Optional[int],
+    device: str,
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    Create a batch dictionary with overlapping frames for model input.
+
+    Args:
+        keys: List of keys to include in the batch
+        value_dict: Dictionary containing values for each key
+        N: Batch dimensions
+        T: Number of frames (optional)
+        device: Device to place tensors on
+
+    Returns:
+        Tuple of (batch dictionary, unconditional batch dictionary)
+    """
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "fps_id":
+            batch[key] = (
+                torch.tensor([value_dict["fps_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "motion_bucket_id":
+            batch[key] = (
+                torch.tensor([value_dict["motion_bucket_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "cond_aug":
+            batch[key] = repeat(
+                torch.tensor([value_dict["cond_aug"]]).to(device),
+                "1 -> b",
+                b=math.prod(N),
+            )
+        elif key == "cond_frames":
+            batch[key] = repeat(value_dict["cond_frames"], "b ... -> (b t) ...", t=N[0])
+        elif key == "cond_frames_without_noise":
+            batch[key] = repeat(
+                value_dict["cond_frames_without_noise"], "b ... -> (b t) ...", t=N[0]
+            )
+        else:
+            batch[key] = value_dict[key]
+
+    if T is not None:
+        batch["num_video_frames"] = T
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+    return batch, batch_uc
+
+
+@torch.inference_mode()
+def sample_keyframes(
+    model_keyframes: Any,
+    audio_list: torch.Tensor,
+    gt_list: torch.Tensor,
+    masks_list: torch.Tensor,
+    condition: torch.Tensor,
+    num_frames: int,
+    fps_id: int,
+    cond_aug: float,
+    device: str,
+    embbedings: Optional[torch.Tensor],
+    force_uc_zero_embeddings: List[str],
+    n_batch_keyframes: int,
+    added_frames: int,
+    strength: float,
+    scale: Optional[Union[float, List[float]]],
+    gt_as_cond: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Sample keyframes using the keyframe generation model.
+
+    Args:
+        model_keyframes: The keyframe generation model
+        audio_list: List of audio embeddings
+        gt_list: List of ground truth frames
+        masks_list: List of masks
+        condition: Conditioning tensor
+        num_frames: Number of frames to generate
+        fps_id: FPS ID
+        cond_aug: Conditioning augmentation factor
+        device: Device to use for computation
+        embbedings: Optional embeddings
+        force_uc_zero_embeddings: List of embeddings to force to zero in unconditional case
+        n_batch_keyframes: Batch size for keyframe generation
+        added_frames: Number of additional frames
+        strength: Strength parameter for sampling
+        scale: Scale parameter for guidance
+        gt_as_cond: Whether to use ground truth as conditioning
+
+    Returns:
+        Tuple of (latent samples, decoded samples)
+    """
+    if scale is not None:
+        model_keyframes.sampler.guider.set_scale(scale)
+    # samples_list = []
+    samples_z_list = []
+    # samples_x_list = []
+
+    for i in range(audio_list.shape[0]):
+        H, W = condition.shape[-2:]
+        assert condition.shape[1] == 3
+        F = 8
+        C = 4
+        shape = (num_frames, C, H // F, W // F)
+
+        audio_cond = audio_list[i].unsqueeze(0)
+
+        value_dict: Dict[str, Any] = {}
+        value_dict["fps_id"] = fps_id
+        value_dict["cond_aug"] = cond_aug
+        value_dict["cond_frames_without_noise"] = condition
+        if embbedings is not None:
+            value_dict["cond_frames"] = embbedings + cond_aug * torch.randn_like(
+                embbedings
+            )
+        else:
+            value_dict["cond_frames"] = condition + cond_aug * torch.randn_like(
+                condition
+            )
+        gt = rearrange(gt_list[i].unsqueeze(0), "b t c h w -> b c t h w").to(device)
+
+        if gt_as_cond:
+            value_dict["cond_frames"] = gt[:, :, 0]
+
+        value_dict["cond_aug"] = cond_aug
+        value_dict["audio_emb"] = audio_cond
+
+        value_dict["gt"] = gt
+        value_dict["masks"] = masks_list[i].unsqueeze(0).transpose(1, 2).to(device)
+
+        with torch.no_grad():
+            batch, batch_uc = get_batch(
+                get_unique_embedder_keys_from_conditioner(model_keyframes.conditioner),
+                value_dict,
+                [1, 1],
+                T=num_frames,
+                device=device,
+            )
+
+            c, uc = model_keyframes.conditioner.get_unconditional_conditioning(
+                batch,
+                batch_uc=batch_uc,
+                force_uc_zero_embeddings=force_uc_zero_embeddings,
+            )
+
+            for k in ["crossattn"]:
+                if c[k].shape[1] != num_frames:
+                    uc[k] = repeat(
+                        uc[k],
+                        "b ... -> b t ...",
+                        t=num_frames,
+                    )
+                    uc[k] = rearrange(
+                        uc[k],
+                        "b t ... -> (b t) ...",
+                        t=num_frames,
+                    )
+                    c[k] = repeat(
+                        c[k],
+                        "b ... -> b t ...",
+                        t=num_frames,
+                    )
+                    c[k] = rearrange(
+                        c[k],
+                        "b t ... -> (b t) ...",
+                        t=num_frames,
+                    )
+
+            video = torch.randn(shape, device=device)
+
+            additional_model_inputs: Dict[str, torch.Tensor] = {}
+            additional_model_inputs["image_only_indicator"] = torch.zeros(
+                n_batch_keyframes, num_frames
+            ).to(device)
+            additional_model_inputs["num_video_frames"] = batch["num_video_frames"]
+
+            def denoiser(
+                input: torch.Tensor, sigma: torch.Tensor, c: Dict[str, torch.Tensor]
+            ) -> torch.Tensor:
+                return model_keyframes.denoiser(
+                    model_keyframes.model,
+                    input,
+                    sigma,
+                    c,
+                    **additional_model_inputs,
+                )
+
+            samples_z = model_keyframes.sampler(
+                denoiser, video, cond=c, uc=uc, strength=strength
+            )
+            samples_z_list.append(samples_z)
+            # samples_x_list.append(samples_x)
+            # samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+            # samples_list.append(samples)
+
+            video = None
+
+    # samples = (
+    #     torch.concat(samples_list)[:-added_frames]
+    #     if added_frames > 0
+    #     else torch.concat(samples_list)
+    # )
+    samples_z = (
+        torch.concat(samples_z_list)[:-added_frames]
+        if added_frames > 0
+        else torch.concat(samples_z_list)
+    )
+    # samples_x = (
+    #     torch.concat(samples_x_list)[:-added_frames]
+    #     if added_frames > 0
+    #     else torch.concat(samples_x_list)
+    # )
+
+    return samples_z
+
+
+@torch.inference_mode()
+def sample_interpolation(
+    model: Any,
+    samples_z: torch.Tensor,
+    # samples_x: torch.Tensor,
+    audio_interpolation_list: List[torch.Tensor],
+    gt_chunks: List[torch.Tensor],
+    masks_chunks: List[torch.Tensor],
+    condition: torch.Tensor,
+    num_frames: int,
+    device: str,
+    overlap: int,
+    fps_id: int,
+    cond_aug: float,
+    force_uc_zero_embeddings: List[str],
+    n_batch: int,
+    chunk_size: Optional[int],
+    strength: float,
+    scale: Optional[float] = None,
+    cut_audio: bool = False,
+    to_remove: List[bool] = [],
+) -> np.ndarray:
+    """
+    Sample interpolation frames between keyframes.
+
+    Args:
+        model: The interpolation model
+        samples_z: Latent samples from keyframe generation
+        samples_x: Decoded samples from keyframe generation
+        audio_interpolation_list: List of audio embeddings for interpolation
+        gt_chunks: Ground truth video chunks
+        masks_chunks: Mask chunks for conditional generation
+        condition: Visual conditioning
+        num_frames: Number of frames to generate
+        device: Device to run inference on
+        overlap: Number of frames to overlap between segments
+        fps_id: FPS ID for conditioning
+        motion_bucket_id: Motion bucket ID for conditioning
+        cond_aug: Conditioning augmentation strength
+        force_uc_zero_embeddings: Keys to zero out in unconditional embeddings
+        n_batch: Batch size for generation
+        chunk_size: Size of chunks for processing (to manage memory)
+        strength: Strength of the conditioning
+        scale: Optional scale for classifier-free guidance
+        cut_audio: Whether to cut audio embeddings
+        to_remove: List of flags indicating which frames to remove
+
+    Returns:
+        Generated video frames as numpy array
+    """
+    if scale is not None:
+        model.sampler.guider.set_scale(scale)
+
+    # Creating condition for interpolation model. We need to create a list of inputs, each input is  [first, last]
+    # The first and last are the first and last frames of the interpolation
+    # interpolation_cond_list = []
+    interpolation_cond_list_emb = []
+
+    # samples_x = [sample for i, sample in zip(to_remove, samples_x) if not i]
+    samples_z = [sample for i, sample in zip(to_remove, samples_z) if not i]
+
+    for i in range(0, len(samples_z) - 1, overlap if overlap > 0 else 2):
+        # interpolation_cond_list.append(
+        #     torch.stack([samples_x[i], samples_x[i + 1]], dim=1)
+        # )
+        interpolation_cond_list_emb.append(
+            torch.stack([samples_z[i], samples_z[i + 1]], dim=1)
+        )
+
+    # condition = torch.stack(interpolation_cond_list).to(device)
+    audio_cond = torch.stack(audio_interpolation_list).to(device)
+    embbedings = torch.stack(interpolation_cond_list_emb).to(device)
+
+    gt_chunks = torch.stack(gt_chunks).to(device)
+    masks_chunks = torch.stack(masks_chunks).to(device)
+
+    H, W = 512, 512
+    F = 8
+    C = 4
+    shape = (num_frames * audio_cond.shape[0], C, H // F, W // F)
+
+    value_dict: Dict[str, Any] = {}
+    value_dict["fps_id"] = fps_id
+    value_dict["cond_aug"] = cond_aug
+    # value_dict["cond_frames_without_noise"] = condition
+
+    value_dict["cond_frames"] = embbedings
+    value_dict["cond_aug"] = cond_aug
+    if cut_audio:
+        value_dict["audio_emb"] = audio_cond[:, :, :, :768]
+    else:
+        value_dict["audio_emb"] = audio_cond
+
+    value_dict["gt"] = rearrange(gt_chunks, "b t c h w -> b c t h w").to(device)
+    value_dict["masks"] = masks_chunks.transpose(1, 2).to(device)
+
+    with torch.no_grad():
+        with torch.autocast(device):
+            batch, batch_uc = get_batch_overlap(
+                get_unique_embedder_keys_from_conditioner(model.conditioner),
+                value_dict,
+                [1, num_frames],
+                T=num_frames,
+                device=device,
+            )
+
+            c, uc = model.conditioner.get_unconditional_conditioning(
+                batch,
+                batch_uc=batch_uc,
+                force_uc_zero_embeddings=force_uc_zero_embeddings,
+            )
+
+            for k in ["crossattn"]:
+                if c[k].shape[1] != num_frames:
+                    uc[k] = repeat(uc[k], "b ... -> b t ...", t=num_frames)
+                    uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=num_frames)
+                    c[k] = repeat(c[k], "b ... -> b t ...", t=num_frames)
+                    c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=num_frames)
+
+            video = torch.randn(shape, device=device)
+
+            additional_model_inputs: Dict[str, torch.Tensor] = {}
+            additional_model_inputs["image_only_indicator"] = torch.zeros(
+                n_batch, num_frames
+            ).to(device)
+            additional_model_inputs["num_video_frames"] = batch["num_video_frames"]
+
+            # Debug information
+            print(
+                f"Shapes - Embeddings: {embbedings.shape}, "
+                f"Audio: {audio_cond.shape}, Video: {shape}, Additional inputs: {additional_model_inputs}"
+            )
+
+            if chunk_size is not None:
+                chunk_size = chunk_size * num_frames
+
+            def denoiser(
+                input: torch.Tensor, sigma: torch.Tensor, c: Dict[str, torch.Tensor]
+            ) -> torch.Tensor:
+                return model.denoiser(
+                    model.model,
+                    input,
+                    sigma,
+                    c,
+                    num_overlap_frames=overlap,
+                    num_frames=num_frames,
+                    n_skips=n_batch,
+                    chunk_size=chunk_size,
+                    **additional_model_inputs,
+                )
+
+            samples_z = model.sampler(denoiser, video, cond=c, uc=uc, strength=strength)
+            samples_z = rearrange(samples_z, "(b t) c h w -> b t c h w", t=num_frames)
+            samples_z[:, 0] = embbedings[:, :, 0]
+            samples_z[:, -1] = embbedings[:, :, 1]
+            samples_z = rearrange(samples_z, "b t c h w -> (b t) c h w")
+
+            samples_x = model.decode_first_stage(samples_z)
+
+            samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+
+            # Free up memory
+            video = None
+
+    samples = rearrange(samples, "(b t) c h w -> b t c h w", t=num_frames)
+    samples = merge_overlapping_segments(samples, overlap)
+
+    vid = (
+        (rearrange(samples, "t c h w -> t c h w") * 255).cpu().numpy().astype(np.uint8)
+    )
+
+    return vid

landmarks_extractor.py

@@ -0,0 +1,35 @@
+from skimage import io
+import face_alignment
+
+
+class LandmarksExtractor:
+    def __init__(self, device="cuda", landmarks_type="2D", flip=False):
+        self.fa = face_alignment.FaceAlignment(
+            face_alignment.LandmarksType.TWO_D
+            if landmarks_type == "2D"
+            else face_alignment.LandmarksType.THREE_D,
+            flip_input=flip,
+            device=device,
+            face_detector="sfd",
+        )
+
+        self.landmarks = []
+
+    def cuda(self):
+        return self
+
+    def extract_landmarks(self, image):
+        # image: either a path to an image or a numpy array (H, W, C) or tensor batch  (B, C, H, W)
+        if isinstance(image, str):
+            image = io.imread(image)
+
+        # Ensure image is on CPU
+        if hasattr(image, "device"):
+            image = image.cpu()
+
+        if len(image.shape) == 3:
+            preds = self.fa.get_landmarks(image)
+        else:
+            preds = self.fa.get_landmarks_from_batch(image)
+
+        return preds

sgm/__init__.py

@@ -0,0 +1,4 @@
+from .models import AutoencodingEngine, DiffusionEngine
+from .util import get_configs_path, instantiate_from_config
+
+__version__ = "0.1.0"

sgm/callbacks/custom_ddp.py

@@ -0,0 +1,10 @@
+# from pytorch_lightning.overrides import LightningDistributedModule
+from pytorch_lightning.strategies import DDPStrategy
+
+
+class CustomDDPPlugin(DDPStrategy):
+    def configure_ddp(self):
+        # self.pre_configure_ddp()
+        self._model = self._setup_model((self.model))
+        self._register_ddp_hooks()
+        self._model._set_static_graph()  # THIS IS THE MAGIC LINE

sgm/callbacks/image_logger.py

@@ -0,0 +1,193 @@
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.loggers import WandbLogger
+import numpy as np
+from pytorch_lightning.utilities import rank_zero_only
+from typing import Union
+import pytorch_lightning as pl
+import os
+from matplotlib import pyplot as plt
+from sgm.util import exists, isheatmap
+import torchvision
+from PIL import Image
+import torch
+import wandb
+from einops import rearrange
+
+
+class ImageLogger(Callback):
+    def __init__(
+        self,
+        batch_frequency,
+        max_images,
+        clamp=True,
+        increase_log_steps=True,
+        rescale=True,
+        disabled=False,
+        log_on_batch_idx=False,
+        log_first_step=False,
+        log_images_kwargs=None,
+        log_before_first_step=False,
+        enable_autocast=True,
+    ):
+        super().__init__()
+        self.enable_autocast = enable_autocast
+        self.rescale = rescale
+        self.batch_freq = batch_frequency
+        self.max_images = max_images
+        self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)]
+        if not increase_log_steps:
+            self.log_steps = [self.batch_freq]
+        self.clamp = clamp
+        self.disabled = disabled
+        self.log_on_batch_idx = log_on_batch_idx
+        self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {}
+        self.log_first_step = log_first_step
+        self.log_before_first_step = log_before_first_step
+
+    @rank_zero_only
+    def log_local(
+        self,
+        save_dir,
+        split,
+        images,
+        global_step,
+        current_epoch,
+        batch_idx,
+        pl_module: Union[None, pl.LightningModule] = None,
+    ):
+        root = os.path.join(save_dir, "images", split)
+        for k in images:
+            if isheatmap(images[k]):
+                fig, ax = plt.subplots()
+                ax = ax.matshow(images[k].cpu().numpy(), cmap="hot", interpolation="lanczos")
+                plt.colorbar(ax)
+                plt.axis("off")
+
+                filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx)
+                os.makedirs(root, exist_ok=True)
+                path = os.path.join(root, filename)
+                plt.savefig(path)
+                plt.close()
+                # TODO: support wandb
+            else:
+                grid = torchvision.utils.make_grid(images[k].squeeze(2), nrow=4)
+                if self.rescale:
+                    grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+                # print(grid.shape, grid.dtype, grid.min(), grid.max(), k)
+                grid = rearrange(grid.squeeze(1), "c h w -> h w c")
+                # grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                grid = grid.numpy()
+                grid = (grid * 255).astype(np.uint8)
+                filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx)
+                path = os.path.join(root, filename)
+                os.makedirs(os.path.split(path)[0], exist_ok=True)
+                img = Image.fromarray(grid)
+                img.save(path)
+                if exists(pl_module):
+                    assert isinstance(
+                        pl_module.logger, WandbLogger
+                    ), "logger_log_image only supports WandbLogger currently"
+                    pl_module.logger.log_image(
+                        key=f"{split}/{k}",
+                        images=[
+                            img,
+                        ],
+                        step=pl_module.global_step,
+                    )
+
+    @rank_zero_only
+    def log_img(self, pl_module, batch, batch_idx, split="train"):
+        check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step
+        if (
+            self.check_frequency(check_idx)
+            and hasattr(pl_module, "log_images")  # batch_idx % self.batch_freq == 0
+            and callable(pl_module.log_images)
+            and
+            # batch_idx > 5 and
+            self.max_images > 0
+        ):
+            logger = type(pl_module.logger)
+            is_train = pl_module.training
+            if is_train:
+                pl_module.eval()
+
+            gpu_autocast_kwargs = {
+                "enabled": self.enable_autocast,  # torch.is_autocast_enabled(),
+                "dtype": torch.get_autocast_gpu_dtype(),
+                "cache_enabled": torch.is_autocast_cache_enabled(),
+            }
+            with torch.no_grad(), torch.cuda.amp.autocast(**gpu_autocast_kwargs):
+                images = pl_module.log_images(batch, split=split, **self.log_images_kwargs)
+
+            for k in images:
+                N = min(images[k].shape[0], self.max_images)
+                if not isheatmap(images[k]):
+                    images[k] = images[k][:N]
+                if isinstance(images[k], torch.Tensor):
+                    images[k] = images[k].detach().float().cpu()
+                    if self.clamp and not isheatmap(images[k]):
+                        images[k] = torch.clamp(images[k], -1.0, 1.0)
+
+            self.log_local(
+                pl_module.logger.save_dir,
+                split,
+                images,
+                pl_module.global_step,
+                pl_module.current_epoch,
+                batch_idx,
+                pl_module=pl_module if isinstance(pl_module.logger, WandbLogger) else None,
+            )
+
+            if is_train:
+                pl_module.train()
+
+    def check_frequency(self, check_idx):
+        if check_idx:
+            check_idx -= 1
+        if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (
+            check_idx > 0 or self.log_first_step
+        ):
+            try:
+                self.log_steps.pop(0)
+            except IndexError as e:
+                print(e)
+                pass
+            return True
+        return False
+
+    @rank_zero_only
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):
+            self.log_img(pl_module, batch, batch_idx, split="train")
+
+    @rank_zero_only
+    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
+        if self.log_before_first_step and pl_module.global_step == 0:
+            print(f"{self.__class__.__name__}: logging before training")
+            self.log_img(pl_module, batch, batch_idx, split="train")
+
+    @rank_zero_only
+    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, *args, **kwargs):
+        if not self.disabled and pl_module.global_step > 0:
+            self.log_img(pl_module, batch, batch_idx, split="val")
+        if hasattr(pl_module, "calibrate_grad_norm"):
+            if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0:
+                self.log_gradients(trainer, pl_module, batch_idx=batch_idx)
+
+
+@rank_zero_only
+def init_wandb(save_dir, opt, config, group_name, name_str):
+    print(f"setting WANDB_DIR to {save_dir}")
+    os.makedirs(save_dir, exist_ok=True)
+
+    os.environ["WANDB_DIR"] = save_dir
+    if opt.debug:
+        wandb.init(project=opt.projectname, mode="offline", group=group_name)
+    else:
+        wandb.init(
+            project=opt.projectname,
+            config=config,
+            settings=wandb.Settings(code_dir="./sgm"),
+            group=group_name,
+            name=name_str,
+        )

sgm/callbacks/setup_callback.py

@@ -0,0 +1,86 @@
+from pytorch_lightning.callbacks import Callback
+import pytorch_lightning as pl
+import os
+from omegaconf import OmegaConf
+from pytorch_lightning.utilities import rank_zero_only
+
+MULTINODE_HACKS = True
+
+
+class SetupCallback(Callback):
+    def __init__(
+        self,
+        resume,
+        now,
+        logdir,
+        ckptdir,
+        cfgdir,
+        config,
+        lightning_config,
+        debug,
+        ckpt_name=None,
+    ):
+        super().__init__()
+        self.resume = resume
+        self.now = now
+        self.logdir = logdir
+        self.ckptdir = ckptdir
+        self.cfgdir = cfgdir
+        self.config = config
+        self.lightning_config = lightning_config
+        self.debug = debug
+        self.ckpt_name = ckpt_name
+
+    @rank_zero_only
+    def on_exception(self, trainer: pl.Trainer, pl_module, exception):
+        print("Exception occurred: {}".format(exception))
+        if not self.debug and trainer.global_rank == 0:
+            print("Summoning checkpoint.")
+            if self.ckpt_name is None:
+                ckpt_path = os.path.join(self.ckptdir, "last.ckpt")
+            else:
+                ckpt_path = os.path.join(self.ckptdir, self.ckpt_name)
+            trainer.save_checkpoint(ckpt_path)
+
+    @rank_zero_only
+    def on_fit_start(self, trainer, pl_module):
+        if trainer.global_rank == 0:
+            # Create logdirs and save configs
+            os.makedirs(self.logdir, exist_ok=True)
+            os.makedirs(self.ckptdir, exist_ok=True)
+            os.makedirs(self.cfgdir, exist_ok=True)
+
+            if "callbacks" in self.lightning_config:
+                if "metrics_over_trainsteps_checkpoint" in self.lightning_config["callbacks"]:
+                    os.makedirs(
+                        os.path.join(self.ckptdir, "trainstep_checkpoints"),
+                        exist_ok=True,
+                    )
+            print("Project config")
+            print(OmegaConf.to_yaml(self.config))
+            if MULTINODE_HACKS:
+                import time
+
+                time.sleep(5)
+            OmegaConf.save(
+                self.config,
+                os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)),
+            )
+
+            print("Lightning config")
+            print(OmegaConf.to_yaml(self.lightning_config))
+            OmegaConf.save(
+                OmegaConf.create({"lightning": self.lightning_config}),
+                os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)),
+            )
+
+        else:
+            # ModelCheckpoint callback created log directory --- remove it
+            if not MULTINODE_HACKS and not self.resume and os.path.exists(self.logdir):
+                dst, name = os.path.split(self.logdir)
+                dst = os.path.join(dst, "child_runs", name)
+                os.makedirs(os.path.split(dst)[0], exist_ok=True)
+                try:
+                    os.rename(self.logdir, dst)
+                except FileNotFoundError:
+                    pass

sgm/callbacks/video_logger.py

@@ -0,0 +1,294 @@
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.loggers import WandbLogger
+import numpy as np
+from pytorch_lightning.utilities import rank_zero_only
+from typing import Union
+import pytorch_lightning as pl
+import os
+from sgm.util import exists, suppress_output, default
+import torchvision
+from PIL import Image
+import torch
+import wandb
+import moviepy.editor as mpy
+from einops import rearrange
+import torchaudio
+# import tempfile
+# import cv2
+# import scipy.io.wavfile as wav
+# import ffmpeg
+
+
+@suppress_output
+def save_audio_video(
+    video, audio=None, frame_rate=25, sample_rate=16000, save_path="temp.mp4", keep_intermediate=False
+):
+    """Save audio and video to a single file.
+    video: (t, c, h, w)
+    audio: (channels t)
+    """
+
+    # temp_filename = next(tempfile._get_candidate_names())
+    # if save_path:
+    #     save_path = save_path
+    # else:
+    #     save_path = "/tmp/" + next(tempfile._get_candidate_names()) + ".mp4"
+    save_path = str(save_path)
+    try:
+        torchvision.io.write_video(
+            "temp_video.mp4", rearrange(video.detach().cpu(), "t c h w -> t h w c").to(torch.uint8), frame_rate
+        )
+        video_clip = mpy.VideoFileClip("temp_video.mp4")
+        if audio is not None:
+            torchaudio.save("temp_audio.wav", audio.detach().cpu(), sample_rate)
+            audio_clip = mpy.AudioFileClip("temp_audio.wav")
+            video_clip = video_clip.set_audio(audio_clip)
+        video_clip.write_videofile(save_path, fps=frame_rate, codec="libx264", audio_codec="aac", verbose=False)
+        if not keep_intermediate:
+            os.remove("temp_video.mp4")
+            if audio is not None:
+                os.remove("temp_audio.wav")
+        return 1
+    except Exception as e:
+        print(e)
+        print("Saving video to file failed")
+        return 0
+
+
+# def write_video_opencv(video, video_rate, video_path):
+#     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+#     out = cv2.VideoWriter(video_path, fourcc, video_rate, (video.shape[2], video.shape[3]), 0)
+#     for frame in list(video):
+#         frame = np.squeeze(frame)
+#         out.write(np.squeeze(frame))
+#     out.release()
+
+
+# # Code mostly inherited from bulletin
+# def save_av_sample(video, video_rate, audio=None, audio_rate=16_000, path=None):
+#     # Save video sample in train dir for debugging
+#     # video_save = 0.5 * video.detach().cpu().numpy() + 0.5
+#     video_save = rearrange(video, "t c h w -> t h w c").detach().cpu().numpy()
+#     temp_filename = next(tempfile._get_candidate_names())
+#     if path:
+#         video_path = path
+#     else:
+#         video_path = "/tmp/" + next(tempfile._get_candidate_names()) + ".mp4"
+#     write_video_opencv((video_save).astype(np.uint8), video_rate, "/tmp/" + temp_filename + ".mp4")
+#     audio_save = audio.detach().squeeze().cpu().numpy()
+#     wav.write("/tmp/" + temp_filename + ".wav", audio_rate, audio_save)
+#     try:
+#         in1 = ffmpeg.input("/tmp/" + temp_filename + ".mp4")
+#         in2 = ffmpeg.input("/tmp/" + temp_filename + ".wav")
+#         out = ffmpeg.output(in1["v"], in2["a"], video_path, loglevel="panic").overwrite_output()
+#         out.run(capture_stdout=True, capture_stderr=True)
+#     except ffmpeg.Error as e:
+#         print("stdout:", e.stdout.decode("utf8"))
+#         print("stderr:", e.stderr.decode("utf8"))
+#         raise e
+#     return video_path
+
+
+class VideoLogger(Callback):
+    def __init__(
+        self,
+        batch_frequency,
+        max_videos,
+        clamp=True,
+        increase_log_steps=True,
+        rescale=True,
+        disabled=False,
+        log_on_batch_idx=False,
+        log_first_step=False,
+        log_videos_kwargs=None,
+        log_before_first_step=False,
+        enable_autocast=True,
+        batch_frequency_val=None,
+    ):
+        super().__init__()
+        self.enable_autocast = enable_autocast
+        self.rescale = rescale
+        self.batch_freq = batch_frequency
+        self.max_videos = max_videos
+        self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)]
+        if not increase_log_steps:
+            self.log_steps = [self.batch_freq]
+        self.batch_freq_val = default(batch_frequency_val, self.batch_freq)
+        self.log_steps_val = [2**n for n in range(int(np.log2(self.batch_freq_val)) + 1)]
+        if not increase_log_steps:
+            self.log_steps_val = [self.batch_freq_val]
+        self.clamp = clamp
+        self.disabled = disabled
+        self.log_on_batch_idx = log_on_batch_idx
+        self.log_videos_kwargs = log_videos_kwargs if log_videos_kwargs else {}
+        self.log_first_step = log_first_step
+        self.log_before_first_step = log_before_first_step
+
+    @rank_zero_only
+    def log_local(
+        self,
+        save_dir,
+        split,
+        log_elements,
+        raw_audio,
+        global_step,
+        current_epoch,
+        batch_idx,
+        pl_module: Union[None, pl.LightningModule] = None,
+    ):
+        root = os.path.join(save_dir, "videos", split)
+        for k in log_elements:
+            element = log_elements[k]
+            if len(element.shape) == 4:
+                grid = torchvision.utils.make_grid(element, nrow=4)
+                if self.rescale:
+                    grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+                grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                grid = grid.numpy()
+                grid = (grid * 255).astype(np.uint8)
+                filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx)
+                path = os.path.join(root, filename)
+                os.makedirs(os.path.split(path)[0], exist_ok=True)
+                img = Image.fromarray(grid)
+                img.save(path)
+                if exists(pl_module):
+                    assert isinstance(
+                        pl_module.logger, WandbLogger
+                    ), "logger_log_image only supports WandbLogger currently"
+                    pl_module.logger.log_image(
+                        key=f"{split}/{k}",
+                        images=[
+                            img,
+                        ],
+                        step=pl_module.global_step,
+                    )
+            elif len(element.shape) == 5:
+                video = element
+                if self.rescale:
+                    video = (video + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+                video = video * 255.0
+                video = video.permute(0, 2, 1, 3, 4).cpu().detach().to(torch.uint8)  # b,t,c,h,w
+                for i in range(video.shape[0]):
+                    filename = "{}_gs-{:06}_e-{:06}_b-{:06}_{}.mp4".format(k, global_step, current_epoch, batch_idx, i)
+                    path = os.path.join(root, filename)
+                    os.makedirs(os.path.split(path)[0], exist_ok=True)
+                    log_audio = raw_audio[i] if raw_audio is not None else None
+                    success = save_audio_video(
+                        video[i],
+                        audio=log_audio.unsqueeze(0) if log_audio is not None else None,
+                        frame_rate=25,
+                        sample_rate=16000,
+                        save_path=path,
+                        keep_intermediate=False,
+                    )
+
+                    # video_path = save_av_sample(video[i], 25, audio=raw_audio, audio_rate=16000, path=None)
+                    if exists(pl_module):
+                        assert isinstance(
+                            pl_module.logger, WandbLogger
+                        ), "logger_log_image only supports WandbLogger currently"
+                        pl_module.logger.experiment.log(
+                            {
+                                f"{split}/{k}": wandb.Video(
+                                    path if success else video,
+                                    # caption=f"diffused videos w {n_frames} frames (condition left, generated right)",
+                                    fps=25,
+                                    format="mp4",
+                                )
+                            },
+                        )
+
+    @rank_zero_only
+    def log_video(self, pl_module, batch, batch_idx, split="train"):
+        check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step
+        # print(f"check_idx: {check_idx}", f"split: {split}")
+        if (
+            self.check_frequency(check_idx, split=split)
+            and hasattr(pl_module, "log_videos")  # batch_idx % self.batch_freq == 0
+            and callable(pl_module.log_videos)
+            and
+            # batch_idx > 5 and
+            self.max_videos > 0
+        ):
+            logger = type(pl_module.logger)
+            is_train = pl_module.training
+            if is_train:
+                pl_module.eval()
+
+            gpu_autocast_kwargs = {
+                "enabled": self.enable_autocast,  # torch.is_autocast_enabled(),
+                "dtype": torch.get_autocast_gpu_dtype(),
+                "cache_enabled": torch.is_autocast_cache_enabled(),
+            }
+            with torch.no_grad(), torch.cuda.amp.autocast(**gpu_autocast_kwargs):
+                videos = pl_module.log_videos(batch, split=split, **self.log_videos_kwargs)
+
+            for k in videos:
+                N = min(videos[k].shape[0], self.max_videos)
+                videos[k] = videos[k][:N]
+                if isinstance(videos[k], torch.Tensor):
+                    videos[k] = videos[k].detach().float().cpu()
+                    if self.clamp:
+                        videos[k] = torch.clamp(videos[k], -1.0, 1.0)
+
+            raw_audio = batch.get("raw_audio", None)
+
+            self.log_local(
+                pl_module.logger.save_dir,
+                split,
+                videos,
+                raw_audio,
+                pl_module.global_step,
+                pl_module.current_epoch,
+                batch_idx,
+                pl_module=pl_module if isinstance(pl_module.logger, WandbLogger) else None,
+            )
+
+            if is_train:
+                pl_module.train()
+
+    def check_frequency(self, check_idx, split="train"):
+        if split == "val":
+            if check_idx:
+                check_idx -= 1
+            if ((check_idx % self.batch_freq_val) == 0 or (check_idx in self.log_steps_val)) and (
+                check_idx > 0 or self.log_first_step
+            ):
+                try:
+                    self.log_steps_val.pop(0)
+                except IndexError as e:
+                    print(e)
+                    pass
+                return True
+            return False
+        if check_idx:
+            check_idx -= 1
+        if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (
+            check_idx > 0 or self.log_first_step
+        ):
+            try:
+                self.log_steps.pop(0)
+            except IndexError as e:
+                print(e)
+                pass
+            return True
+        return False
+
+    @rank_zero_only
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):
+            self.log_video(pl_module, batch, batch_idx, split="train")
+
+    @rank_zero_only
+    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
+        if self.log_before_first_step and pl_module.global_step == 0:
+            print(f"{self.__class__.__name__}: logging before training")
+            self.log_video(pl_module, batch, batch_idx, split="train")
+
+    @rank_zero_only
+    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, *args, **kwargs):
+        if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):
+            self.log_video(pl_module, batch, batch_idx, split="val")
+        if hasattr(pl_module, "calibrate_grad_norm"):
+            if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0:
+                self.log_gradients(trainer, pl_module, batch_idx=batch_idx)

sgm/data/__init__.py

@@ -0,0 +1 @@
+# from .dataset import StableDataModuleFromConfig

sgm/data/data_utils.py

@@ -0,0 +1,561 @@
+import torch
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+from functools import partial
+import math
+
+
+def get_size(img):
+    if isinstance(img, (np.ndarray, torch.Tensor)):
+        return img.shape[1::-1]
+    else:
+        return img.size
+
+
+def imresample(img, sz):
+    im_data = torch.nn.functional.interpolate(img, size=sz, mode="area")
+    return im_data
+
+
+def crop_resize(img, box, image_size):
+    if isinstance(img, np.ndarray):
+        img = img[box[1] : box[3], box[0] : box[2]]
+        out = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_AREA).copy()
+    elif isinstance(img, torch.Tensor):
+        img = img[box[1] : box[3], box[0] : box[2]]
+        out = (
+            imresample(img.permute(2, 0, 1).unsqueeze(0).float(), (image_size, image_size))
+            .byte()
+            .squeeze(0)
+            .permute(1, 2, 0)
+        )
+    else:
+        out = img.crop(box).copy().resize((image_size, image_size), Image.BILINEAR)
+    return out
+
+
+def fixed_image_standardization(image_tensor):
+    processed_tensor = (image_tensor - 127.5) / 128.0
+    return processed_tensor
+
+
+def extract_face(img, landmarks, image_size=160, margin=0, postprocess=False):
+    """Extract face + margin from images given facial landmarks.
+
+    Arguments:
+        img {PIL.Image/torch.Tensor/np.ndarray} -- Input image(s) with shape (B, H, W, C)
+        landmarks {numpy.ndarray} -- Facial landmarks with shape (B, 68, 2)
+        image_size {int} -- Output image size in pixels. The image will be square.
+        margin {int} -- Margin to add to bounding box, in terms of pixels in the final image.
+        postprocess {bool} -- Whether to apply standardization
+
+    Returns:
+        torch.tensor -- tensor representing the extracted faces with shape (B, H, W, C)
+    """
+    # Calculate bounding boxes from landmarks for all faces in batch
+    x_min = np.min(landmarks, axis=1)[:, 0]  # Shape: (B,)
+    y_min = np.min(landmarks, axis=1)[:, 1]  # Shape: (B,)
+    x_max = np.max(landmarks, axis=1)[:, 0]  # Shape: (B,)
+    y_max = np.max(landmarks, axis=1)[:, 1]  # Shape: (B,)
+
+    # Calculate margin for top only
+    box_height = y_max - y_min
+    top_margin = margin * box_height / (image_size - margin)
+
+    # Create boxes for all faces
+    boxes = np.stack(
+        [
+            x_min,
+            np.maximum(y_min - top_margin, 0),  # Only add margin to top
+            x_max,
+            y_max,
+        ],
+        axis=1,
+    ).astype(int)  # Shape: (B, 4)
+
+    # Process each face in the batch
+    faces = []
+    for i in range(len(boxes)):
+        face = crop_resize(img[i], boxes[i], image_size)
+        faces.append(face)
+
+    faces = torch.stack(faces, dim=0)
+    faces = faces.float()
+
+    if postprocess:
+        faces = fixed_image_standardization(faces)
+
+    return faces
+
+
+def crop_mouth_region(images, landmarks, crop_size=96):
+    """
+    Takes a fixed-size square crop centered on the mouth region.
+
+    Parameters:
+    - images: tensor/array of shape (num_frames, height, width, channels) or (height, width, channels)
+    - landmarks: numpy array of shape (num_frames, 68, 2) or (68, 2)
+    - crop_size: size of the square crop (both height and width)
+    - padding: percentage of padding around the mouth region (0.0 to 1.0)
+
+    Returns:
+    - List of fixed-size crops or single crop if input is single image
+    """
+    # Handle single image case
+    single_image = False
+    if len(images.shape) == 3:
+        images = images[None]
+        landmarks = landmarks[None]
+        single_image = True
+
+    num_frames = len(images)
+    crops = []
+
+    # Mouth landmarks indices (48-67 for mouth region)
+    mouth_indices = range(48, 68)
+
+    for i in range(num_frames):
+        # Get mouth landmarks for current frame
+        mouth_landmarks = landmarks[i][mouth_indices]
+
+        # Find center of mouth
+        center_x = int(np.mean(mouth_landmarks[:, 0]))
+        center_y = int(np.mean(mouth_landmarks[:, 1]))
+
+        # Calculate crop boundaries
+        half_size = crop_size // 2
+        left = max(0, center_x - half_size)
+        right = min(images.shape[2], center_x + half_size)
+        top = max(0, center_y - half_size)
+        bottom = min(images.shape[1], center_y + half_size)
+
+        # Adjust if crop would go out of bounds
+        if left == 0:
+            right = crop_size
+        if right == images.shape[2]:
+            left = images.shape[2] - crop_size
+        if top == 0:
+            bottom = crop_size
+        if bottom == images.shape[1]:
+            top = images.shape[1] - crop_size
+
+        # Take the crop
+        crop = images[i, top:bottom, left:right]
+        crops.append(crop)
+
+    return crops[0] if single_image else crops
+
+
+def create_masks_from_landmarks_box(landmark_list, img_shape, nose_index=28, dtype="uint8", box_expand=0.0):
+    height, width = img_shape[:2]
+    num_frames = landmark_list.shape[0]
+
+    # Initialize the masks array
+    masks = np.zeros((num_frames, height, width), dtype=dtype)
+
+    if 0 <= box_expand < 1:
+        box_expand = int(box_expand * width)
+
+    for i in range(num_frames):
+        # Get the landmarks for the current frame
+        landmarks = landmark_list[i]
+
+        # Get the y-coordinate of the nose landmark
+        nose_point_h = landmarks[nose_index, 1]
+        cut_h = nose_point_h
+
+        # Find the leftmost and rightmost landmarks
+        far_left_index = np.argmin(landmarks[:, 0])
+        far_right_index = np.argmax(landmarks[:, 0])
+
+        # Define the points for the mask contour
+        left_up_point = np.array([landmarks[far_left_index][0], cut_h - box_expand], dtype=np.int32)
+        left_down_point = np.array([landmarks[far_left_index][0], height], dtype=np.int32)
+        right_up_point = np.array([landmarks[far_right_index][0], cut_h - box_expand], dtype=np.int32)
+        right_down_point = np.array([landmarks[far_right_index][0], height], dtype=np.int32)
+
+        # Define the contour
+        contour = np.array([[left_up_point, left_down_point, right_down_point, right_up_point]])
+
+        # Draw the contour on the mask
+        cv2.drawContours(masks[i], [contour], -1, color=(1), thickness=cv2.FILLED)
+
+    return torch.from_numpy(masks)
+
+
+def create_masks_from_landmarks_full_size(
+    landmarks_batch, image_height, image_width, start_index=48, end_index=68, offset=0, nose_index=33
+):
+    """
+    Efficiently creates a batch of masks using vectorized operations where each mask has ones from the highest
+    landmark in the specified range (adjusted by an offset) to the bottom of the image, and zeros otherwise.
+
+    Parameters:
+    - landmarks_batch (np.array): An array of shape (B, 68, 2) containing facial landmarks for multiple samples.
+    - image_height (int): The height of the image for which masks are created.
+    - image_width (int): The width of the image for which masks are created.
+    - start_index (int): The starting index of the range to check (inclusive).
+    - end_index (int): The ending index of the range to check (inclusive).
+    - offset (int): An offset to add or subtract from the y-coordinate of the highest landmark.
+
+    Returns:
+    - np.array: An array of masks of shape (B, image_height, image_width) for each batch.
+    """
+    # Extract the y-coordinates for the specified range across all batches
+    y_coords = landmarks_batch[:, nose_index : nose_index + 1, 1]
+
+    # Find the index of the minimum y-coordinate in the specified range for each batch
+    min_y_indices = np.argmin(y_coords, axis=1)
+
+    # Gather the highest landmarks' y-coordinates using the indices found
+    highest_y_coords = y_coords[np.arange(len(y_coords)), min_y_indices]
+
+    if abs(offset) < 1 and abs(offset) > 0:
+        offset = int(offset * image_height)
+
+    # Apply the offset to the highest y-coordinate
+    adjusted_y_coords = highest_y_coords + offset
+
+    # Clip the coordinates to stay within image boundaries
+    adjusted_y_coords = np.clip(adjusted_y_coords, 0, image_height - 1)
+
+    # Use broadcasting to create a mask without loops
+    # Create a range of indices from 0 to image_height - 1
+    all_indices = np.arange(image_height)
+
+    # Compare each index in 'all_indices' to each 'adjusted_y_coord' in the batch
+    # 'all_indices' has shape (image_height,), we reshape to (1, image_height) to broadcast against (B, 1)
+    mask_2d = (all_indices >= adjusted_y_coords[:, None]).astype(int)
+
+    # Extend the 2D mask to a full 3D mask of size (B, image_height, image_width)
+    full_mask = np.tile(mask_2d[:, :, np.newaxis], (1, 1, image_width))
+
+    return torch.from_numpy(full_mask)
+
+
+def expand_polygon(polygon, expand_size):
+    """
+    Expands the polygon outward by a specified number of pixels.
+
+    Parameters:
+    - polygon (list of tuples): The polygon points as (x, y).
+    - expand_size (int): The number of pixels to expand the polygon outward.
+
+    Returns:
+    - expanded_polygon (list of tuples): The expanded polygon points as (x, y).
+    """
+    if expand_size == 0:
+        return polygon
+
+    # Calculate centroid of the polygon
+    centroid_x = sum([point[0] for point in polygon]) / len(polygon)
+    centroid_y = sum([point[1] for point in polygon]) / len(polygon)
+
+    # Expand each point outward from the centroid
+    expanded_polygon = []
+    for x, y in polygon:
+        vector_x = x - centroid_x
+        vector_y = y - centroid_y
+        length = np.sqrt(vector_x**2 + vector_y**2)
+        if length == 0:
+            expanded_polygon.append((x, y))
+        else:
+            new_x = x + expand_size * (vector_x / length)
+            new_y = y + expand_size * (vector_y / length)
+            expanded_polygon.append((int(new_x), int(new_y)))
+
+    return expanded_polygon
+
+
+def create_masks_from_landmarks_mouth(landmark_list, img_shape, nose_index=33, dtype="uint8", box_expand=0.0):
+    height, width = img_shape[:2]
+    num_frames = landmark_list.shape[0]
+
+    # Initialize the masks array
+    masks = np.zeros((num_frames, height, width), dtype=dtype)
+
+    if 0 <= box_expand < 1:
+        box_expand = int(box_expand * width)
+
+    for i in range(num_frames):
+        # Get the landmarks for the current frame
+        landmarks = landmark_list[i]
+
+        # Get the y-coordinate of the nose landmark
+        nose_point_h = landmarks[nose_index, 1]
+        cut_h = nose_point_h
+
+        # Find the leftmost and rightmost landmarks
+        far_left_index = np.argmin(landmarks[:, 0])
+        far_right_index = np.argmax(landmarks[:, 0])
+
+        # Find lowest landmark y-coordinate
+        lowest_y = np.max(landmarks[:, 1])
+        # Add box_expand to the lowest point
+        lowest_y = min(height, lowest_y + box_expand)
+
+        # Define the points for the mask contour
+        left_up_point = np.array([landmarks[far_left_index][0], cut_h - box_expand], dtype=np.int32)
+        left_down_point = np.array([landmarks[far_left_index][0], lowest_y], dtype=np.int32)
+        right_up_point = np.array([landmarks[far_right_index][0], cut_h - box_expand], dtype=np.int32)
+        right_down_point = np.array([landmarks[far_right_index][0], lowest_y], dtype=np.int32)
+
+        # Define the contour
+        contour = np.array([[left_up_point, left_down_point, right_down_point, right_up_point]])
+
+        # Draw the contour on the mask
+        cv2.drawContours(masks[i], [contour], -1, color=(1), thickness=cv2.FILLED)
+
+    return torch.from_numpy(masks)
+
+
+def create_face_mask_from_landmarks(landmarks_batch, image_height, image_width, mask_expand=0):
+    """
+    Creates a batch of masks where each mask covers the face region using landmarks.
+
+    Parameters:
+    - landmarks_batch (np.array): An array of shape (B, 68, 2) containing facial landmarks for multiple samples.
+    - image_height (int): The height of the image for which masks are created.
+    - image_width (int): The width of the image for which masks are created.
+    - mask_expand (int): The number of pixels to expand the mask outward.
+
+    Returns:
+    - np.array: An array of masks of shape (B, image_height, image_width) for each batch.
+    """
+    # Initialize an array to hold all masks
+    masks = np.zeros((landmarks_batch.shape[0], image_height, image_width), dtype=np.uint8)
+
+    if abs(mask_expand) < 1 and abs(mask_expand) > 0:
+        mask_expand = int(mask_expand * image_height)
+
+    for i, landmarks in enumerate(landmarks_batch):
+        # Create a blank image for each mask
+        mask = Image.new("L", (image_width, image_height), 0)
+        draw = ImageDraw.Draw(mask)
+
+        # Extract relevant landmarks for the face
+        jawline_landmarks = landmarks[2:15]  # Jawline
+        # upper_face_landmarks = landmarks[17:27]  # Eyebrows and top of nose bridge
+
+        # Combine landmarks to form a polygon around the face
+        # face_polygon = np.concatenate((jawline_landmarks, upper_face_landmarks[::-1]), axis=0)
+        face_polygon = jawline_landmarks
+
+        # Convert landmarks to a list of tuples
+        face_polygon = [(int(x), int(y)) for x, y in face_polygon]
+
+        # Expand the polygon if necessary
+        expanded_polygon = expand_polygon(face_polygon, mask_expand)
+
+        # Draw the polygon and fill it
+        draw.polygon(expanded_polygon, outline=1, fill=1)
+
+        # Convert mask to numpy array and add it to the batch of masks
+        masks[i] = np.array(mask)
+
+    return torch.from_numpy(masks)
+
+
+ALL_FIXED_POINTS = (
+    [i for i in range(0, 4)] + [i for i in range(13, 17)] + [i for i in range(27, 36)] + [36, 39, 42, 45]
+)
+
+
+def gaussian_kernel(sigma, width, height):
+    """Create a 2D Gaussian kernel."""
+    x = torch.arange(0, width, 1) - width // 2
+    y = torch.arange(0, height, 1) - height // 2
+    x = x.float()
+    y = y.float()
+    x2 = x**2
+    y2 = y[:, None] ** 2
+    g = torch.exp(-(x2 + y2) / (2 * sigma**2))
+    return g / g.sum()
+
+
+def generate_hm(landmarks, height, width, n_points="all", sigma=3):
+    if n_points == "all":
+        Nlandmarks = range(len(landmarks))
+    elif n_points == "fixed":
+        Nlandmarks = ALL_FIXED_POINTS
+    elif n_points == "stable":
+        Nlandmarks = [33, 36, 39, 42, 45]
+
+    kernel = gaussian_kernel(sigma, width, height)
+    hm = torch.zeros((height, width))
+    for I in Nlandmarks:
+        x0, y0 = landmarks[I]
+        x0, y0 = int(x0), int(y0)
+        left, right = max(0, x0 - width // 2), min(width, x0 + width // 2)
+        top, bottom = max(0, y0 - height // 2), min(height, y0 + height // 2)
+        hm[top:bottom, left:right] += kernel[
+            max(0, -y0 + height // 2) : min(height, height - y0 + height // 2),
+            max(0, -x0 + width // 2) : min(width, width - x0 + width // 2),
+        ]
+    # Normalize the heatmap to have values between 0 and 1
+    max_val = hm.max()
+    if max_val > 0:
+        hm /= max_val
+    return hm
+
+
+def get_heatmap(landmarks, image_size, or_im_size, n_points="stable", sigma=4):
+    stack = []
+    seq_length = landmarks.shape[0]
+    if or_im_size[0] != image_size[0] or or_im_size[1] != image_size[1]:
+        landmarks = scale_landmarks(landmarks, or_im_size, image_size)
+    gen_single_heatmap = partial(
+        generate_hm,
+        height=image_size[0],
+        width=image_size[1],
+        n_points=n_points,
+        sigma=sigma,
+    )
+    for i in range(seq_length):
+        stack.append(gen_single_heatmap(landmarks[i]))
+
+    return torch.stack(stack, axis=0).unsqueeze(0)  # (1, seq_length, height, width)
+
+
+def scale_landmarks(landmarks, original_size, target_size):
+    """
+    Scale landmarks from original size to target size.
+
+    Parameters:
+    - landmarks (np.array): An array of shape (N, 2) containing facial landmarks.
+    - original_size (tuple): The size (height, width) for which the landmarks are currently scaled.
+    - target_size (tuple): The size (height, width) to which landmarks should be scaled.
+
+    Returns:
+    - scaled_landmarks (np.array): Scaled landmarks.
+    """
+    scale_y = target_size[0] / original_size[0]
+    scale_x = target_size[1] / original_size[1]
+    scaled_landmarks = landmarks * np.array([scale_x, scale_y])
+    return scaled_landmarks.astype(int)
+
+
+def draw_kps_image(
+    image_shape, original_size, landmarks, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255)], rgb=True, pts_width=4
+):
+    stick_width = pts_width
+    limb_seq = np.array([[0, 2], [1, 2]])
+    kps = landmarks[[36, 45, 33], :]
+    kps = scale_landmarks(kps, original_size, image_shape)
+    if not rgb:  # Grayscale image
+        canvas = np.zeros((image_shape[0], image_shape[1], 1))
+        color_mode = "grayscale"
+    else:  # Color image
+        canvas = np.zeros((image_shape[0], image_shape[1], 3))
+        color_mode = "color"
+
+    polygon_cache = {}
+
+    for index in limb_seq:
+        color = color_list[index[0]]
+        if color_mode == "grayscale":
+            color = (int(0.299 * color[2] + 0.587 * color[1] + 0.114 * color[0]),)  # Convert to grayscale intensity
+
+        x = kps[index][:, 0]
+        y = kps[index][:, 1]
+        length = np.sqrt((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2)
+        angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1]))
+
+        cache_key = (color, int(np.mean(x)), int(np.mean(y)), int(length / 2), int(angle))
+        if cache_key not in polygon_cache:
+            polygon_cache[cache_key] = cv2.ellipse2Poly(
+                (int(np.mean(x)), int(np.mean(y))), (int(length / 2), stick_width), int(angle), 0, 360, 1
+            )
+
+        polygon = polygon_cache[cache_key]
+        cv2.fillConvexPoly(canvas, polygon, [int(c * 0.6) for c in color])
+
+    for idx, kp in enumerate(kps):
+        if color_mode == "grayscale":
+            color = (int(0.299 * color_list[idx][2] + 0.587 * color_list[idx][1] + 0.114 * color_list[idx][0]),)
+        else:
+            color = color_list[idx]
+        cv2.circle(canvas, (int(kp[0]), int(kp[1])), pts_width, color, -1)
+
+    return canvas.transpose(2, 0, 1)
+
+
+def create_landmarks_image(
+    landmarks, original_size=(772, 772), target_size=(772, 772), point_size=3, n_points="all", dim=3
+):
+    """
+    Creates an image of landmarks on a black background using efficient NumPy operations.
+
+    Parameters:
+    - landmarks (np.array): An array of shape (68, 2) containing facial landmarks.
+    - image_size (tuple): The size of the output image (height, width).
+    - point_size (int): The radius of each landmark point in pixels.
+
+    Returns:
+    - img (np.array): An image array with landmarks plotted.
+    """
+    if n_points == "all":
+        indexes = range(len(landmarks))
+    elif n_points == "fixed":
+        indexes = ALL_FIXED_POINTS
+    elif n_points == "stable":
+        indexes = [33, 36, 39, 42, 45]
+
+    landmarks = landmarks[indexes]
+
+    img = np.zeros(target_size, dtype=np.uint8)
+
+    landmarks = scale_landmarks(landmarks, original_size, target_size)
+
+    # Ensure the landmarks are in bounds and integer
+    landmarks = np.clip(landmarks, [0, 0], [target_size[1] - 1, target_size[0] - 1]).astype(int)
+
+    # Get x and y coordinates from landmarks
+    x, y = landmarks[:, 0], landmarks[:, 1]
+
+    # Define a grid offset based on point_size around each landmark
+    offset = np.arange(-point_size // 2, point_size // 2 + 1)
+    grid_x, grid_y = np.meshgrid(offset, offset, indexing="ij")
+
+    # Calculate the full set of x and y coordinates for the points
+    full_x = x[:, np.newaxis, np.newaxis] + grid_x[np.newaxis, :, :]
+    full_y = y[:, np.newaxis, np.newaxis] + grid_y[np.newaxis, :, :]
+
+    # Clip the coordinates to stay within image boundaries
+    full_x = np.clip(full_x, 0, target_size[1] - 1)
+    full_y = np.clip(full_y, 0, target_size[0] - 1)
+
+    # Flatten the arrays to use them as indices
+    full_x = full_x.ravel()
+    full_y = full_y.ravel()
+
+    # Set the points in the image
+    img[full_y, full_x] = 255
+
+    return np.stack([img] * dim, axis=0)
+
+
+def trim_pad_audio(audio, sr, max_len_sec=None, max_len_raw=None):
+    len_file = audio.shape[-1]
+
+    if max_len_sec or max_len_raw:
+        max_len = max_len_raw if max_len_raw is not None else int(max_len_sec * sr)
+        if len_file < int(max_len):
+            # dummy = np.zeros((1, int(max_len_sec * sr) - len_file))
+            # extened_wav = np.concatenate((audio_data, dummy[0]))
+            extened_wav = torch.nn.functional.pad(audio, (0, int(max_len) - len_file), "constant")
+        else:
+            extened_wav = audio[:, : int(max_len)]
+    else:
+        extened_wav = audio
+
+    return extened_wav
+
+
+def ssim_to_bin(ssim_score):
+    # Normalize the SSIM score to a 0-100 scale
+    normalized_diff_ssim = (1 - ((ssim_score + 1) / 2)) * 100
+    # Assign to one of the 100 bins
+    bin_index = float(min(np.floor(normalized_diff_ssim), 99))
+    return bin_index

sgm/data/dataset.py

@@ -0,0 +1,80 @@
+from typing import Optional
+
+import torchdata.datapipes.iter
+import webdataset as wds
+from omegaconf import DictConfig
+from pytorch_lightning import LightningDataModule
+
+try:
+    from sdata import create_dataset, create_dummy_dataset, create_loader
+except ImportError as e:
+    print("#" * 100)
+    print("Datasets not yet available")
+    print("to enable, we need to add stable-datasets as a submodule")
+    print("please use ``git submodule update --init --recursive``")
+    print("and do ``pip install -e stable-datasets/`` from the root of this repo")
+    print("#" * 100)
+    exit(1)
+
+
+class StableDataModuleFromConfig(LightningDataModule):
+    def __init__(
+        self,
+        train: DictConfig,
+        validation: Optional[DictConfig] = None,
+        test: Optional[DictConfig] = None,
+        skip_val_loader: bool = False,
+        dummy: bool = False,
+    ):
+        super().__init__()
+        self.train_config = train
+        assert (
+            "datapipeline" in self.train_config and "loader" in self.train_config
+        ), "train config requires the fields `datapipeline` and `loader`"
+
+        self.val_config = validation
+        if not skip_val_loader:
+            if self.val_config is not None:
+                assert (
+                    "datapipeline" in self.val_config and "loader" in self.val_config
+                ), "validation config requires the fields `datapipeline` and `loader`"
+            else:
+                print(
+                    "Warning: No Validation datapipeline defined, using that one from training"
+                )
+                self.val_config = train
+
+        self.test_config = test
+        if self.test_config is not None:
+            assert (
+                "datapipeline" in self.test_config and "loader" in self.test_config
+            ), "test config requires the fields `datapipeline` and `loader`"
+
+        self.dummy = dummy
+        if self.dummy:
+            print("#" * 100)
+            print("USING DUMMY DATASET: HOPE YOU'RE DEBUGGING ;)")
+            print("#" * 100)
+
+    def setup(self, stage: str) -> None:
+        print("Preparing datasets")
+        if self.dummy:
+            data_fn = create_dummy_dataset
+        else:
+            data_fn = create_dataset
+
+        self.train_datapipeline = data_fn(**self.train_config.datapipeline)
+        if self.val_config:
+            self.val_datapipeline = data_fn(**self.val_config.datapipeline)
+        if self.test_config:
+            self.test_datapipeline = data_fn(**self.test_config.datapipeline)
+
+    def train_dataloader(self) -> torchdata.datapipes.iter.IterDataPipe:
+        loader = create_loader(self.train_datapipeline, **self.train_config.loader)
+        return loader
+
+    def val_dataloader(self) -> wds.DataPipeline:
+        return create_loader(self.val_datapipeline, **self.val_config.loader)
+
+    def test_dataloader(self) -> wds.DataPipeline:
+        return create_loader(self.test_datapipeline, **self.test_config.loader)

sgm/data/mask.py

@@ -0,0 +1,525 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+"""
+Functions taken from https://github.com/DanBigioi/DiffusionVideoEditing
+
+
+"""
+
+import cv2
+import numpy as np
+import torch
+
+" Countour from 2:15 not good for head poses "
+
+
+def face_mask(img_shape, landmark_list, dtype="uint8"):
+    height, width = img_shape[:2]
+    mask = np.ones((height, width, 1), dtype=dtype)
+    cv2.drawContours(
+        mask, np.int32([landmark_list[2:15]]), -1, color=(0), thickness=cv2.FILLED
+    )
+
+    return mask
+
+
+def face_mask_jaw_box(img_shape, landmark_list, dtype="uint8", kernel_size=10):
+    nose = 33
+    jaw = 8
+
+    height, width = img_shape[:2]
+    mask = np.ones((height, width, 1), dtype=dtype)
+    combined_landmarks = np.concatenate((landmark_list[2:15], [landmark_list[33]]))
+
+    # Draw the combined contour on the mask
+    cv2.drawContours(
+        mask, [np.int32(combined_landmarks)], -1, color=(0), thickness=cv2.FILLED
+    )
+
+    inverted_mask = 1 - mask
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    mask = cv2.dilate(inverted_mask, kernel, iterations=1)
+    mask = np.expand_dims(
+        mask, axis=-1
+    )  # Add a singleton dimension to match the number of channels
+    mask = 1 - mask
+
+    cut_h = landmark_list[nose][1]
+
+    far_left = int(np.argmin(landmark_list[:, 0]))
+    far_right = int(np.argmax(landmark_list[:, 0]))
+    left_up_point = np.int32([landmark_list[far_left][0], cut_h])  # 2
+    right_up_point = np.int32([landmark_list[far_right][0], cut_h])  # 15
+    height_landmarks = min(landmark_list[jaw, 1] + 20, height)
+    left_down_point = np.int32([landmark_list[far_left][0], height_landmarks])
+    right_down_point = np.int32([landmark_list[far_right][0], height_landmarks])
+
+    # print(cut_h, cut_h + 10, height_landmarks)
+
+    mask_box = [left_up_point, left_down_point, right_down_point, right_up_point]
+
+    return mask, mask_box
+
+
+" Stretch the tight face mask - Countour from 2:15 but dilate, not good for extreme head poses "
+
+
+def face_mask_stretch(img_shape, landmark_list, dtype="uint8", kernel_size=10):
+    height, width = img_shape[:2]
+    mask = np.ones((height, width, 1), dtype=dtype)
+    combined_landmarks = np.concatenate((landmark_list[2:15], [landmark_list[33]]))
+
+    # Draw the combined contour on the mask
+    cv2.drawContours(
+        mask, [np.int32(combined_landmarks)], -1, color=(0), thickness=cv2.FILLED
+    )
+
+    # cv2.drawContours(mask, np.int32([landmark_list[2:15]]), -1, color=(0), thickness=cv2.FILLED)
+    inverted_mask = 1 - mask
+
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    mask = cv2.dilate(inverted_mask, kernel, iterations=1)
+    mask = np.expand_dims(
+        mask, axis=-1
+    )  # Add a singleton dimension to match the number of channels
+    mask = 1 - mask
+
+    return mask
+
+
+" Small box around mouth - Use far left, far right points for extreme head poses, cut between nose and upper mouth point"
+
+
+def face_mask_box_pose(img_shape, landmark_list, dtype="uint8"):
+    """
+    When the head pose is different than frontal then the normal cropping with landmarks does not work correctly.
+    Crop using as height the middle nose point
+    Take the left/right corners using the far_left and far_right landmarks
+    TODO: Maybe it is better to add some more pixels to have a bigger mask, especially on large head poses
+    """
+
+    height, width = img_shape[:2]
+
+    nose = 33
+    upper_lip = 51
+    jaw = 8
+
+    nose_point_h = landmark_list[nose, 1]
+    upper_lip_point = landmark_list[upper_lip, 1]
+    cut_h = (upper_lip_point - nose_point_h) / 2 + nose_point_h
+
+    # cut_h = landmark_list[nose][1]
+
+    mask = np.ones((height, width, 1), dtype=dtype)
+
+    far_left = int(np.argmin(landmark_list[:, 0]))
+    far_right = int(np.argmax(landmark_list[:, 0]))
+
+    left_up_point = np.int32([landmark_list[far_left][0], cut_h])  # 2
+    right_up_point = np.int32([landmark_list[far_right][0], cut_h])  # 15
+
+    height_landmarks = min(landmark_list[jaw, 1] + 20, height)
+    left_down_point = np.int32([landmark_list[far_left][0], height_landmarks])
+    right_down_point = np.int32([landmark_list[far_right][0], height_landmarks])
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    left_up_point,
+                    left_down_point,
+                    right_up_point,
+                    right_down_point,
+                    left_up_point,
+                    right_up_point,
+                    left_down_point,
+                    right_down_point,
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    return mask
+
+
+" Small box around mouth - Use far left, far right points for extreme head poses, cut from nose"
+
+
+def face_mask_box_pose_nose(
+    img_shape,
+    landmark_list,
+    dtype="uint8",
+    get_box=False,
+    pixels_above_nose=None,
+    pixels_under_jaw=None,
+):
+    height, width = img_shape[:2]
+
+    nose = 33
+    jaw = 8
+
+    cut_h = landmark_list[nose][1]
+    if pixels_above_nose is not None:
+        # this is only for inference to take a bigger mask and blend it back to the original frame
+        cut_h = cut_h - pixels_above_nose
+
+    mask = np.ones((height, width, 1), dtype=dtype)
+
+    far_left = int(np.argmin(landmark_list[:, 0]))
+    far_right = int(np.argmax(landmark_list[:, 0]))
+
+    left_up_point = np.int32([landmark_list[far_left][0], cut_h])  # 2
+    right_up_point = np.int32([landmark_list[far_right][0], cut_h])  # 15
+
+    height_landmarks = min(landmark_list[jaw, 1] + 20, height)
+    if pixels_under_jaw is not None:
+        height_landmarks = min(landmark_list[jaw, 1] + pixels_under_jaw, height)
+    left_down_point = np.int32([landmark_list[far_left][0], height_landmarks])
+    right_down_point = np.int32([landmark_list[far_right][0], height_landmarks])
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    left_up_point,
+                    left_down_point,
+                    right_up_point,
+                    right_down_point,
+                    left_up_point,
+                    right_up_point,
+                    left_down_point,
+                    right_down_point,
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    if get_box:
+        mask_box = [left_up_point, left_down_point, right_down_point, right_up_point]
+        return mask, mask_box
+    else:
+        return mask
+
+
+def face_mask_box_pose_big(
+    img_shape, landmark_list, dtype="uint8", cut_h=None, far_left=None, far_right=None
+):
+    height, width = img_shape[:2]
+    mask = np.ones((height, width, 1), dtype=dtype)
+    nose = 33
+    nose_point_h = landmark_list[nose, 1]
+    if cut_h is None:
+        cut_h = nose_point_h
+
+    if far_right is None and far_left is None:
+        far_left = int(np.argmin(landmark_list[:, 0]))
+        far_right = int(np.argmax(landmark_list[:, 0]))
+
+        left_up_point = np.int32([landmark_list[far_left][0], cut_h])
+        left_down_point = np.int32([landmark_list[far_left][0], height])
+
+        right_up_point = np.int32([landmark_list[far_right][0], cut_h])
+        right_down_point = np.int32([landmark_list[far_right][0], height])
+    else:
+        left_up_point = np.int32([far_left, cut_h])
+        left_down_point = np.int32([far_left, height])
+
+        right_up_point = np.int32([far_right, cut_h])
+        right_down_point = np.int32([far_right, height])
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    left_up_point,
+                    left_down_point,
+                    right_up_point,
+                    right_down_point,
+                    left_up_point,
+                    right_up_point,
+                    left_down_point,
+                    right_down_point,
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    return mask
+
+
+def face_mask_box_pose_big_cover_nose(img_shape, landmark_list, dtype="uint8"):
+    height, width = img_shape[:2]
+
+    middle_nose_point = 29
+
+    cut_h = landmark_list[middle_nose_point, 1]
+
+    mask = np.ones((height, width, 1), dtype=dtype)
+
+    far_left = int(np.argmin(landmark_list[:, 0]))
+    far_right = int(np.argmax(landmark_list[:, 0]))
+
+    left_up_point = np.int32([landmark_list[far_left][0], cut_h])
+    left_down_point = np.int32([landmark_list[far_left][0], height])
+
+    right_up_point = np.int32([landmark_list[far_right][0], cut_h])
+    right_down_point = np.int32([landmark_list[far_right][0], height])
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    left_up_point,
+                    left_down_point,
+                    right_up_point,
+                    right_down_point,
+                    left_up_point,
+                    right_up_point,
+                    left_down_point,
+                    right_down_point,
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    return mask
+
+
+def face_mask_square(img_shape, landmark_list, dtype="uint8"):
+    height, width = img_shape[:2]
+
+    mask = np.ones((height, width, 1), dtype=dtype)
+
+    far_left = np.min(landmark_list[:, 0])
+    far_right = np.max(landmark_list[:, 1])
+    print("far_left {}, far_right {}".format(far_left, far_right))
+
+    left_p = 2
+    right_p = 14
+
+    print(
+        "left_p {}, right_p {}".format(
+            landmark_list[left_p][0], landmark_list[right_p][0]
+        )
+    )
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    landmark_list[left_p],
+                    [landmark_list[left_p][0], height],
+                    landmark_list[right_p],
+                    [landmark_list[right_p][0], height],
+                    landmark_list[left_p],
+                    landmark_list[right_p],
+                    [landmark_list[left_p][0], height],
+                    [landmark_list[right_p][0], height],
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    return mask
+
+
+" Used for half face "
+
+
+def bbox2mask(img_shape, bbox, dtype="uint8"):
+    """Generate mask in ndarray from bbox.
+
+    The returned mask has the shape of (h, w, 1). '1' indicates the
+    hole and '0' indicates the valid regions.
+
+    We prefer to use `uint8` as the data type of masks, which may be different
+    from other codes in the community.
+
+    Args:
+            img_shape (tuple[int]): The size of the image.
+            bbox (tuple[int]): Configuration tuple, (top, left, height, width)
+            dtype (str): Indicate the data type of returned masks. Default: 'uint8'
+
+    Return:
+            numpy.ndarray: Mask in the shape of (h, w, 1).
+    """
+
+    height, width = img_shape[:2]
+
+    mask = np.ones((height, width, 1), dtype=dtype)
+    mask[bbox[0] : bbox[0] + bbox[2], bbox[1] : bbox[1] + bbox[3], :] = 0.0
+
+    return mask
+
+
+def face_mask_cheeks(img_shape, landmark_list, dtype="uint8"):
+    height, width = img_shape[:2]
+    mask = np.ones((height, width, 1), dtype=dtype)
+
+    middle_nose_point = 29
+    nose = 33
+    cut_h = int(landmark_list[middle_nose_point, 1])
+
+    far_left = int(np.argmin(landmark_list[:, 0]))
+    far_right = int(np.argmax(landmark_list[:, 0]))
+
+    left_up_point = np.int32([landmark_list[far_left][0], cut_h])
+    left_down_point = np.int32([landmark_list[far_left][0], height])
+
+    right_up_point = np.int32([landmark_list[far_right][0], cut_h])
+    right_down_point = np.int32([landmark_list[far_right][0], height])
+
+    cv2.drawContours(
+        mask,
+        np.int32(
+            [
+                [
+                    left_up_point,
+                    left_down_point,
+                    right_up_point,
+                    right_down_point,
+                    left_up_point,
+                    right_up_point,
+                    left_down_point,
+                    right_down_point,
+                ]
+            ]
+        ),
+        -1,
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    #  Calculate the bounding box coordinates for the nose
+    nose_jaw_dist = (
+        abs(landmark_list[2][0] - landmark_list[middle_nose_point][0]) * 0.10
+    )  # 1, 15
+    # nose_right_dist = (landmark_list[middle_nose_point][0] - landmark_list[1][0]) * 0.10
+    # nose_left_dist = (landmark_list[15][0] - landmark_list[middle_nose_point][0]) * 0.10
+    #
+
+    nose_min_x = int(landmark_list[31][0] - nose_jaw_dist)
+    nose_max_x = int(landmark_list[35][0] + nose_jaw_dist)
+    # nose_min_x = int(landmark_list[31][0] - nose_right_dist)
+    # nose_max_x = int(landmark_list[35][0] + nose_left_dist)
+    nose_min_y = cut_h
+    nose_max_y = int(landmark_list[nose, 1])
+
+    # Clear the nose area from the mask using a rectangle
+    mask_nose = np.ones((height, width, 1), dtype=dtype)
+    cv2.rectangle(
+        mask_nose,
+        (nose_min_x, nose_min_y),
+        (nose_max_x, nose_max_y),
+        color=(0),
+        thickness=cv2.FILLED,
+    )
+
+    mask_nose = 1 - mask_nose
+    mask = mask + mask_nose
+
+    return mask
+
+
+def face_mask_cheeks_batch(
+    img_shape, landmark_list, dtype="uint8", box_expand=0.0, show_nose=True
+):
+    height, width = img_shape[:2]
+
+    # Handle both single and multiple landmarks
+    if len(landmark_list.shape) == 2:
+        landmark_list = landmark_list[None, ...]  # Add batch dimension
+    num_frames = landmark_list.shape[0]
+
+    # Initialize masks for all frames
+    masks = np.ones((num_frames, height, width), dtype=dtype)
+
+    for i in range(num_frames):
+        landmarks = landmark_list[i]
+        middle_nose_point = 29
+        nose = 33
+        cut_h = int(landmarks[middle_nose_point, 1])
+
+        # Add height expansion
+        if box_expand > 0:
+            cut_h = max(0, cut_h - int(box_expand * height))
+
+        far_left = int(np.argmin(landmarks[:, 0]))
+        far_right = int(np.argmax(landmarks[:, 0]))
+
+        left_up_point = np.int32([landmarks[far_left][0], cut_h])
+        left_down_point = np.int32([landmarks[far_left][0], height])
+
+        right_up_point = np.int32([landmarks[far_right][0], cut_h])
+        right_down_point = np.int32([landmarks[far_right][0], height])
+
+        cv2.drawContours(
+            masks[i],
+            np.int32(
+                [
+                    [
+                        left_up_point,
+                        left_down_point,
+                        right_up_point,
+                        right_down_point,
+                        left_up_point,
+                        right_up_point,
+                        left_down_point,
+                        right_down_point,
+                    ]
+                ]
+            ),
+            -1,
+            color=(0),
+            thickness=cv2.FILLED,
+        )
+
+        if show_nose:
+            #  Calculate the bounding box coordinates for the nose
+            nose_jaw_dist = (
+                abs(landmarks[2][0] - landmarks[middle_nose_point][0]) * 0.10
+            )  # 1, 15
+
+            nose_min_x = int(landmarks[31][0] - nose_jaw_dist)
+            nose_max_x = int(landmarks[35][0] + nose_jaw_dist)
+            nose_min_y = cut_h
+            nose_max_y = int(landmarks[nose, 1])
+
+            # Clear the nose area from the mask using a rectangle
+            mask_nose = np.ones((height, width), dtype=dtype)
+            cv2.rectangle(
+                mask_nose,
+                (nose_min_x, nose_min_y),
+                (nose_max_x, nose_max_y),
+                color=(0),
+                thickness=cv2.FILLED,
+            )
+
+            mask_nose = 1 - mask_nose
+            masks[i] = masks[i] + mask_nose
+
+    # If input was single frame, return single mask
+    if landmark_list.shape[0] == 1:
+        return masks[0]
+
+    return 1 - torch.from_numpy(masks)

sgm/data/video_datamodule_latent.py

@@ -0,0 +1,138 @@
+from typing import Any, Dict, Optional
+
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
+from omegaconf import DictConfig
+
+import sys
+import pyrootutils
+
+root = pyrootutils.setup_root(__file__, pythonpath=True)
+sys.path.append(root)
+from sgm.data.video_dataset_latent import VideoDataset
+
+
+class VideoDataModule(LightningDataModule):
+    """
+    A DataModule implements 5 key methods:
+
+        def prepare_data(self):
+            # things to do on 1 GPU/TPU (not on every GPU/TPU in DDP)
+            # download data, pre-process, split, save to disk, etc...
+        def setup(self, stage):
+            # things to do on every process in DDP
+            # load data, set variables, etc...
+        def train_dataloader(self):
+            # return train dataloader
+        def val_dataloader(self):
+            # return validation dataloader
+        def test_dataloader(self):
+            # return test dataloader
+        def teardown(self):
+            # called on every process in DDP
+            # clean up after fit or test
+
+    This allows you to share a full dataset without explaining how to download,
+    split, transform and process the data.
+
+    Read the docs:
+        https://pytorch-lightning.readthedocs.io/en/latest/data/datamodule.html
+    """
+
+    def __init__(
+        self,
+        train: DictConfig,
+        validation: Optional[DictConfig] = None,
+        test: Optional[DictConfig] = None,
+        skip_val_loader: bool = False,
+    ):
+        super().__init__()
+
+        # this line allows to access init params with 'self.hparams' attribute
+        # also ensures init params will be stored in ckpt
+        self.train_config = train
+        assert "datapipeline" in self.train_config and "loader" in self.train_config, (
+            "train config requires the fields `datapipeline` and `loader`"
+        )
+
+        self.val_config = validation
+        if not skip_val_loader:
+            if self.val_config is not None:
+                assert (
+                    "datapipeline" in self.val_config and "loader" in self.val_config
+                ), "validation config requires the fields `datapipeline` and `loader`"
+            else:
+                print(
+                    "Warning: No Validation datapipeline defined, using that one from training"
+                )
+                self.val_config = train
+
+        self.test_config = test
+        if self.test_config is not None:
+            assert (
+                "datapipeline" in self.test_config and "loader" in self.test_config
+            ), "test config requires the fields `datapipeline` and `loader`"
+
+    def setup(self, stage: Optional[str] = None):
+        """Load data. Set variables: `self.data_train`, `self.data_val`, `self.data_test`.
+
+        This method is called by lightning with both `trainer.fit()` and `trainer.test()`, so be
+        careful not to execute things like random split twice!
+        """
+        print("Preparing datasets")
+
+        self.train_datapipeline = VideoDataset(**self.train_config.datapipeline)
+        if self.val_config:
+            self.val_datapipeline = VideoDataset(**self.val_config.datapipeline)
+        if self.test_config:
+            self.test_datapipeline = VideoDataset(**self.test_config.datapipeline)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_datapipeline, **self.train_config.loader)
+
+    def val_dataloader(self):
+        if self.val_datapipeline:
+            return DataLoader(self.val_datapipeline, **self.val_config.loader)
+        else:
+            return None
+
+    def test_dataloader(self):
+        if self.test_datapipeline:
+            return DataLoader(self.test_datapipeline, **self.test_config.loader)
+        else:
+            return None
+
+    def teardown(self, stage: Optional[str] = None):
+        """Clean up after fit or test."""
+        pass
+
+    def state_dict(self):
+        """Extra things to save to checkpoint."""
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        """Things to do when loading checkpoint."""
+        pass
+
+
+if __name__ == "__main__":
+    import hydra
+    import omegaconf
+    import pyrootutils
+    import cv2
+
+    root = pyrootutils.setup_root(__file__, pythonpath=True)
+    cfg = omegaconf.OmegaConf.load(
+        root / "configs" / "datamodule" / "image_datamodule.yaml"
+    )
+    # cfg.data_dir = str(root / "data")
+    data = hydra.utils.instantiate(cfg)
+    data.prepare_data()
+    data.setup()
+    print(data.data_train.__getitem__(0)[0].shape)
+    batch = next(iter(data.train_dataloader()))
+    identity, target = batch
+    image_identity = (identity[0].permute(1, 2, 0).numpy() + 1) / 2 * 255
+    image_other = (target[0].permute(1, 2, 0).numpy() + 1) / 2 * 255
+    cv2.imwrite("image_identity.png", image_identity[:, :, ::-1])
+    cv2.imwrite("image_other.png", image_other[:, :, ::-1])

sgm/data/video_dataset_latent.py

@@ -0,0 +1,780 @@
+import random
+import numpy as np
+from functools import partial
+from torch.utils.data import Dataset, WeightedRandomSampler
+import torch.nn.functional as F
+import torch
+import math
+import decord
+from einops import rearrange
+from more_itertools import sliding_window
+from omegaconf import ListConfig
+import torchaudio
+import soundfile as sf
+from torchvision.transforms import RandomHorizontalFlip
+from audiomentations import Compose, AddGaussianNoise, PitchShift
+from safetensors.torch import load_file
+from tqdm import tqdm
+import cv2
+from sgm.data.data_utils import (
+    create_masks_from_landmarks_full_size,
+    create_face_mask_from_landmarks,
+    create_masks_from_landmarks_box,
+    create_masks_from_landmarks_mouth,
+)
+from sgm.data.mask import face_mask_cheeks_batch
+
+torchaudio.set_audio_backend("sox_io")
+decord.bridge.set_bridge("torch")
+
+
+def exists(x):
+    return x is not None
+
+
+def trim_pad_audio(audio, sr, max_len_sec=None, max_len_raw=None):
+    len_file = audio.shape[-1]
+
+    if max_len_sec or max_len_raw:
+        max_len = max_len_raw if max_len_raw is not None else int(max_len_sec * sr)
+        if len_file < int(max_len):
+            extened_wav = torch.nn.functional.pad(
+                audio, (0, int(max_len) - len_file), "constant"
+            )
+        else:
+            extened_wav = audio[:, : int(max_len)]
+    else:
+        extened_wav = audio
+
+    return extened_wav
+
+
+# Similar to regular video dataset but trades flexibility for speed
+class VideoDataset(Dataset):
+    def __init__(
+        self,
+        filelist,
+        resize_size=None,
+        audio_folder="Audio",
+        video_folder="CroppedVideos",
+        emotions_folder="emotions",
+        landmarks_folder=None,
+        audio_emb_folder=None,
+        video_extension=".avi",
+        audio_extension=".wav",
+        audio_rate=16000,
+        latent_folder=None,
+        audio_in_video=False,
+        fps=25,
+        num_frames=5,
+        need_cond=True,
+        step=1,
+        mode="prediction",
+        scale_audio=False,
+        augment=False,
+        augment_audio=False,
+        use_latent=False,
+        latent_type="stable",
+        latent_scale=1,  # For backwards compatibility
+        from_audio_embedding=False,
+        load_all_possible_indexes=False,
+        audio_emb_type="wavlm",
+        cond_noise=[-3.0, 0.5],
+        motion_id=255.0,
+        data_mean=None,
+        data_std=None,
+        use_latent_condition=False,
+        skip_frames=0,
+        get_separate_id=False,
+        virtual_increase=1,
+        filter_by_length=False,
+        select_randomly=False,
+        balance_datasets=True,
+        use_emotions=False,
+        get_original_frames=False,
+        add_extra_audio_emb=False,
+        expand_box=0.0,
+        nose_index=28,
+        what_mask="full",
+        get_masks=False,
+    ):
+        self.audio_folder = audio_folder
+        self.from_audio_embedding = from_audio_embedding
+        self.audio_emb_type = audio_emb_type
+        self.cond_noise = cond_noise
+        self.latent_condition = use_latent_condition
+        precomputed_latent = latent_type
+        self.audio_emb_folder = (
+            audio_emb_folder if audio_emb_folder is not None else audio_folder
+        )
+        self.skip_frames = skip_frames
+        self.get_separate_id = get_separate_id
+        self.fps = fps
+        self.virtual_increase = virtual_increase
+        self.select_randomly = select_randomly
+        self.use_emotions = use_emotions
+        self.emotions_folder = emotions_folder
+        self.get_original_frames = get_original_frames
+        self.add_extra_audio_emb = add_extra_audio_emb
+        self.expand_box = expand_box
+        self.nose_index = nose_index
+        self.landmarks_folder = landmarks_folder
+        self.what_mask = what_mask
+        self.get_masks = get_masks
+
+        assert not (exists(data_mean) ^ exists(data_std)), (
+            "Both data_mean and data_std should be provided"
+        )
+
+        if data_mean is not None:
+            data_mean = rearrange(torch.as_tensor(data_mean), "c -> c () () ()")
+            data_std = rearrange(torch.as_tensor(data_std), "c -> c () () ()")
+        self.data_mean = data_mean
+        self.data_std = data_std
+        self.motion_id = motion_id
+        self.latent_folder = (
+            latent_folder if latent_folder is not None else video_folder
+        )
+        self.audio_in_video = audio_in_video
+
+        self.filelist = []
+        self.audio_filelist = []
+        self.landmark_filelist = [] if get_masks else None
+        with open(filelist, "r") as files:
+            for f in files.readlines():
+                f = f.rstrip()
+
+                audio_path = f.replace(video_folder, audio_folder).replace(
+                    video_extension, audio_extension
+                )
+
+                self.filelist += [f]
+                self.audio_filelist += [audio_path]
+                if self.get_masks:
+                    landmark_path = f.replace(video_folder, landmarks_folder).replace(
+                        video_extension, ".npy"
+                    )
+                    self.landmark_filelist += [landmark_path]
+
+        self.resize_size = resize_size
+        if use_latent and not precomputed_latent:
+            self.resize_size *= 4 if latent_type in ["stable", "ldm"] else 8
+        self.scale_audio = scale_audio
+        self.step = step
+        self.use_latent = use_latent
+        self.precomputed_latent = precomputed_latent
+        self.latent_type = latent_type
+        self.latent_scale = latent_scale
+        self.video_ext = video_extension
+        self.video_folder = video_folder
+
+        self.augment = augment
+        self.maybe_augment = RandomHorizontalFlip(p=0.5) if augment else lambda x: x
+        self.maybe_augment_audio = (
+            Compose(
+                [
+                    AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.002, p=0.25),
+                    # TimeStretch(min_rate=0.8, max_rate=1.25, p=0.3),
+                    PitchShift(min_semitones=-1, max_semitones=1, p=0.25),
+                    # Shift(min_fraction=-0.5, max_fraction=0.5, p=0.333),
+                ]
+            )
+            if augment_audio
+            else lambda x, sample_rate: x
+        )
+        self.maybe_augment_audio = partial(
+            self.maybe_augment_audio, sample_rate=audio_rate
+        )
+
+        self.mode = mode
+        if mode == "interpolation":
+            need_cond = False  # Interpolation does not need condition as first and last frame becomes the condition
+        self.need_cond = need_cond  # If need cond will extract one more frame than the number of frames
+        if get_separate_id:
+            self.need_cond = True
+        # It is used for the conditional model when the condition is not on the temporal dimension
+        num_frames = num_frames if not self.need_cond else num_frames + 1
+
+        vr = decord.VideoReader(self.filelist[0])
+        self.video_rate = math.ceil(vr.get_avg_fps())
+        print(f"Video rate: {self.video_rate}")
+        self.audio_rate = audio_rate
+        a2v_ratio = fps / float(self.audio_rate)
+        self.samples_per_frame = math.ceil(1 / a2v_ratio)
+
+        if get_separate_id:
+            assert mode == "prediction", (
+                "Separate identity frame is only supported for prediction mode"
+            )
+            # No need for extra frame if we are getting a separate identity frame
+            self.need_cond = True
+            num_frames -= 1
+        self.num_frames = num_frames
+        self.load_all_possible_indexes = load_all_possible_indexes
+        if load_all_possible_indexes:
+            self._indexes = self._get_indexes(
+                self.filelist, self.audio_filelist, self.landmark_filelist
+            )
+        else:
+            if filter_by_length:
+                self._indexes = self.filter_by_length(
+                    self.filelist, self.audio_filelist, self.landmark_filelist
+                )
+            else:
+                if self.get_masks:
+                    self._indexes = list(
+                        zip(self.filelist, self.audio_filelist, self.landmark_filelist)
+                    )
+                else:
+                    self._indexes = list(
+                        zip(
+                            self.filelist,
+                            self.audio_filelist,
+                            [None] * len(self.filelist),
+                        )
+                    )
+
+        self.balance_datasets = balance_datasets
+        if self.balance_datasets:
+            self.weights = self._calculate_weights()
+            self.sampler = WeightedRandomSampler(
+                self.weights, num_samples=len(self._indexes), replacement=True
+            )
+
+    def __len__(self):
+        return len(self._indexes) * self.virtual_increase
+
+    def _load_landmarks(self, filename, original_size, target_size, indexes):
+        landmarks = np.load(filename, allow_pickle=True)[indexes, :]
+        if self.what_mask == "full":
+            mask = create_masks_from_landmarks_full_size(
+                landmarks,
+                original_size[0],
+                original_size[1],
+                offset=self.expand_box,
+                nose_index=self.nose_index,
+            )
+        elif self.what_mask == "box":
+            mask = create_masks_from_landmarks_box(
+                landmarks,
+                (original_size[0], original_size[1]),
+                box_expand=self.expand_box,
+                nose_index=self.nose_index,
+            )
+        elif self.what_mask == "heart":
+            mask = face_mask_cheeks_batch(
+                original_size, landmarks, box_expand=0.0, show_nose=True
+            )
+        elif self.what_mask == "mouth":
+            mask = create_masks_from_landmarks_mouth(
+                landmarks,
+                (original_size[0], original_size[1]),
+                box_expand=0.01,
+                nose_index=self.nose_index,
+            )
+        else:
+            mask = create_face_mask_from_landmarks(
+                landmarks, original_size[0], original_size[1], mask_expand=0.05
+            )
+        # Interpolate the mask to the target size
+        mask = F.interpolate(
+            mask.unsqueeze(1).float(), size=target_size, mode="nearest"
+        )
+
+        return mask, landmarks
+
+    def get_emotions(self, video_file, video_indexes):
+        emotions_path = video_file.replace(
+            self.video_folder, self.emotions_folder
+        ).replace(self.video_ext, ".pt")
+        emotions = torch.load(emotions_path)
+        return (
+            emotions["valence"][video_indexes],
+            emotions["arousal"][video_indexes],
+            emotions["labels"][video_indexes],
+        )
+
+    def get_frame_indices(self, total_video_frames, select_randomly=False, start_idx=0):
+        if select_randomly:
+            # Randomly select self.num_frames indices from the available range
+            available_indices = list(range(start_idx, total_video_frames))
+            if len(available_indices) < self.num_frames:
+                raise ValueError(
+                    "Not enough frames in the video to sample with given parameters."
+                )
+            indexes = random.sample(available_indices, self.num_frames)
+            return sorted(indexes)  # Sort to maintain temporal order
+        else:
+            # Calculate the maximum possible start index
+            max_start_idx = total_video_frames - (
+                (self.num_frames - 1) * (self.skip_frames + 1) + 1
+            )
+
+            # Generate a random start index
+            if max_start_idx > 0:
+                start_idx = np.random.randint(start_idx, max_start_idx)
+            else:
+                raise ValueError(
+                    "Not enough frames in the video to sample with given parameters."
+                )
+
+            # Generate the indices
+            indexes = [
+                start_idx + i * (self.skip_frames + 1) for i in range(self.num_frames)
+            ]
+
+        return indexes
+
+    def _load_audio(self, filename, max_len_sec, start=None, indexes=None):
+        audio, sr = sf.read(
+            filename,
+            start=math.ceil(start * self.audio_rate),
+            frames=math.ceil(self.audio_rate * max_len_sec),
+            always_2d=True,
+        )  # e.g (16000, 1)
+        audio = audio.T  # (1, 16000)
+        assert sr == self.audio_rate, (
+            f"Audio rate is {sr} but should be {self.audio_rate}"
+        )
+        audio = audio.mean(0, keepdims=True)
+        audio = self.maybe_augment_audio(audio)
+        audio = torch.from_numpy(audio).float()
+        # audio = torchaudio.functional.resample(audio, orig_freq=sr, new_freq=self.audio_rate)
+        audio = trim_pad_audio(audio, self.audio_rate, max_len_sec=max_len_sec)
+        return audio[0]
+
+    def ensure_shape(self, tensors):
+        target_length = self.samples_per_frame
+        processed_tensors = []
+        for tensor in tensors:
+            current_length = tensor.shape[1]
+            diff = current_length - target_length
+            assert abs(diff) <= 5, (
+                f"Expected shape {target_length}, but got {current_length}"
+            )
+            if diff < 0:
+                # Calculate how much padding is needed
+                padding_needed = target_length - current_length
+                # Pad the tensor
+                padded_tensor = F.pad(tensor, (0, padding_needed))
+                processed_tensors.append(padded_tensor)
+            elif diff > 0:
+                # Trim the tensor
+                trimmed_tensor = tensor[:, :target_length]
+                processed_tensors.append(trimmed_tensor)
+            else:
+                # If it's already the correct size
+                processed_tensors.append(tensor)
+        return torch.cat(processed_tensors)
+
+    def normalize_latents(self, latents):
+        if self.data_mean is not None:
+            # Normalize latents to 0 mean and 0.5 std
+            latents = ((latents - self.data_mean) / self.data_std) * 0.5
+        return latents
+
+    def convert_indexes(self, indexes_25fps, fps_from=25, fps_to=60):
+        ratio = fps_to / fps_from
+        indexes_60fps = [int(index * ratio) for index in indexes_25fps]
+        return indexes_60fps
+
+    def _get_frames_and_audio(self, idx):
+        if self.load_all_possible_indexes:
+            indexes, video_file, audio_file, land_file = self._indexes[idx]
+            if self.audio_in_video:
+                vr = decord.AVReader(video_file, sample_rate=self.audio_rate)
+            else:
+                vr = decord.VideoReader(video_file)
+            len_video = len(vr)
+            if "AA_processed" in video_file or "1000actors_nsv" in video_file:
+                len_video *= 25 / 60
+                len_video = int(len_video)
+        else:
+            video_file, audio_file, land_file = self._indexes[idx]
+            if self.audio_in_video:
+                vr = decord.AVReader(video_file, sample_rate=self.audio_rate)
+            else:
+                vr = decord.VideoReader(video_file)
+            len_video = len(vr)
+            if "AA_processed" in video_file or "1000actors_nsv" in video_file:
+                len_video *= 25 / 60
+                len_video = int(len_video)
+
+            indexes = self.get_frame_indices(
+                len_video,
+                select_randomly=self.select_randomly,
+                start_idx=120 if "1000actors_nsv" in video_file else 0,
+            )
+
+        if self.get_separate_id:
+            id_idx = np.random.randint(0, len_video)
+            indexes.insert(0, id_idx)
+
+        if "AA_processed" in video_file or "1000actors_nsv" in video_file:
+            video_indexes = self.convert_indexes(indexes, fps_from=25, fps_to=60)
+            audio_file = audio_file.replace("_output_output", "")
+            if self.audio_emb_type == "wav2vec2" and "AA_processed" in video_file:
+                audio_path_extra = ".safetensors"
+            else:
+                audio_path_extra = f"_{self.audio_emb_type}_emb.safetensors"
+
+            video_path_extra = f"_{self.latent_type}_512_latent.safetensors"
+            audio_path_extra_extra = (
+                ".pt" if "AA_processed" in video_file else "_beats_emb.pt"
+            )
+
+        else:
+            video_indexes = indexes
+            audio_path_extra = f"_{self.audio_emb_type}_emb.safetensors"
+            video_path_extra = f"_{self.latent_type}_512_latent.safetensors"
+            audio_path_extra_extra = "_beats_emb.pt"
+
+        emotions = None
+        if self.use_emotions:
+            emotions = self.get_emotions(video_file, video_indexes)
+            if self.get_separate_id:
+                emotions = (emotions[0][1:], emotions[1][1:], emotions[2][1:])
+
+        raw_audio = None
+        if self.audio_in_video:
+            raw_audio, frames_video = vr.get_batch(video_indexes)
+            raw_audio = rearrange(self.ensure_shape(raw_audio), "f s -> (f s)")
+
+        if self.use_latent and self.precomputed_latent:
+            latent_file = video_file.replace(self.video_ext, video_path_extra).replace(
+                self.video_folder, self.latent_folder
+            )
+            frames = load_file(latent_file)["latents"][video_indexes, :, :, :]
+
+            if frames.shape[-1] != 64:
+                print(f"Frames shape: {frames.shape}, video file: {video_file}")
+
+            frames = rearrange(frames, "t c h w -> c t h w") * self.latent_scale
+            frames = self.normalize_latents(frames)
+        else:
+            if self.audio_in_video:
+                frames = frames_video.permute(3, 0, 1, 2).float()
+            else:
+                frames = vr.get_batch(video_indexes).permute(3, 0, 1, 2).float()
+
+        if raw_audio is None:
+            # Audio is not in video
+            raw_audio = self._load_audio(
+                audio_file,
+                max_len_sec=frames.shape[1] / self.fps,
+                start=indexes[0] / self.fps,
+                # indexes=indexes,
+            )
+        if not self.from_audio_embedding:
+            audio = raw_audio
+            audio_frames = rearrange(audio, "(f s) -> f s", s=self.samples_per_frame)
+        else:
+            audio = load_file(
+                audio_file.replace(self.audio_folder, self.audio_emb_folder).split(".")[
+                    0
+                ]
+                + audio_path_extra
+            )["audio"]
+            audio_frames = audio[indexes, :]
+            if self.add_extra_audio_emb:
+                audio_extra = torch.load(
+                    audio_file.replace(self.audio_folder, self.audio_emb_folder).split(
+                        "."
+                    )[0]
+                    + audio_path_extra_extra
+                )
+                audio_extra = audio_extra[indexes, :]
+                audio_frames = torch.cat([audio_frames, audio_extra], dim=-1)
+
+        audio_frames = (
+            audio_frames[1:] if self.need_cond else audio_frames
+        )  # Remove audio of first frame
+
+        if self.get_original_frames:
+            original_frames = vr.get_batch(video_indexes).permute(3, 0, 1, 2).float()
+            original_frames = self.scale_and_crop((original_frames / 255.0) * 2 - 1)
+            original_frames = (
+                original_frames[:, 1:] if self.need_cond else original_frames
+            )
+        else:
+            original_frames = None
+
+        if not self.use_latent or (self.use_latent and not self.precomputed_latent):
+            frames = self.scale_and_crop((frames / 255.0) * 2 - 1)
+
+        target = frames[:, 1:] if self.need_cond else frames
+        if self.mode == "prediction":
+            if self.use_latent:
+                if self.audio_in_video:
+                    clean_cond = (
+                        frames_video[0].unsqueeze(0).permute(3, 0, 1, 2).float()
+                    )
+                else:
+                    clean_cond = (
+                        vr[video_indexes[0]].unsqueeze(0).permute(3, 0, 1, 2).float()
+                    )
+                original_size = clean_cond.shape[-2:]
+                clean_cond = self.scale_and_crop((clean_cond / 255.0) * 2 - 1).squeeze(
+                    0
+                )
+                if self.latent_condition:
+                    noisy_cond = frames[:, 0]
+                else:
+                    noisy_cond = clean_cond
+            else:
+                clean_cond = frames[:, 0]
+                noisy_cond = clean_cond
+        elif self.mode == "interpolation":
+            if self.use_latent:
+                if self.audio_in_video:
+                    clean_cond = frames_video[[0, -1]].permute(3, 0, 1, 2).float()
+                else:
+                    clean_cond = (
+                        vr.get_batch([video_indexes[0], video_indexes[-1]])
+                        .permute(3, 0, 1, 2)
+                        .float()
+                    )
+                original_size = clean_cond.shape[-2:]
+                clean_cond = self.scale_and_crop((clean_cond / 255.0) * 2 - 1)
+                if self.latent_condition:
+                    noisy_cond = torch.stack([target[:, 0], target[:, -1]], dim=1)
+                else:
+                    noisy_cond = clean_cond
+            else:
+                clean_cond = torch.stack([target[:, 0], target[:, -1]], dim=1)
+                noisy_cond = clean_cond
+
+        # Add noise to conditional frame
+        if self.cond_noise and isinstance(self.cond_noise, ListConfig):
+            cond_noise = (
+                self.cond_noise[0] + self.cond_noise[1] * torch.randn((1,))
+            ).exp()
+            noisy_cond = noisy_cond + cond_noise * torch.randn_like(noisy_cond)
+        else:
+            noisy_cond = noisy_cond + self.cond_noise * torch.randn_like(noisy_cond)
+            cond_noise = self.cond_noise
+
+        if self.get_masks:
+            target_size = (
+                (self.resize_size, self.resize_size)
+                if not self.use_latent
+                else (self.resize_size // 8, self.resize_size // 8)
+            )
+            masks, landmarks = self._load_landmarks(
+                land_file, original_size, target_size, video_indexes
+            )
+
+            landmarks = None
+            masks = (
+                masks.permute(1, 0, 2, 3)[:, 1:]
+                if self.need_cond
+                else masks.permute(1, 0, 2, 3)
+            )
+        else:
+            masks = None
+            landmarks = None
+
+        return (
+            original_frames,
+            clean_cond,
+            noisy_cond,
+            target,
+            audio_frames,
+            raw_audio,
+            cond_noise,
+            emotions,
+            masks,
+            landmarks,
+        )
+
+    def filter_by_length(self, video_filelist, audio_filelist):
+        def with_opencv(filename):
+            video = cv2.VideoCapture(filename)
+            frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT)
+
+            return int(frame_count)
+
+        filtered_video = []
+        filtered_audio = []
+        min_length = (self.num_frames - 1) * (self.skip_frames + 1) + 1
+        for vid_file, audio_file in tqdm(
+            zip(video_filelist, audio_filelist),
+            total=len(video_filelist),
+            desc="Filtering",
+        ):
+            # vr = decord.VideoReader(vid_file)
+
+            len_video = with_opencv(vid_file)
+            # Short videos
+            if len_video < min_length:
+                continue
+            filtered_video.append(vid_file)
+            filtered_audio.append(audio_file)
+        print(f"New number of files: {len(filtered_video)}")
+        return filtered_video, filtered_audio
+
+    def _get_indexes(self, video_filelist, audio_filelist):
+        indexes = []
+        self.og_shape = None
+        for vid_file, audio_file in zip(video_filelist, audio_filelist):
+            vr = decord.VideoReader(vid_file)
+            if self.og_shape is None:
+                self.og_shape = vr[0].shape[-2]
+            len_video = len(vr)
+            # Short videos
+            if len_video < self.num_frames:
+                continue
+            else:
+                possible_indexes = list(
+                    sliding_window(range(len_video), self.num_frames)
+                )[:: self.step]
+                possible_indexes = list(
+                    map(lambda x: (x, vid_file, audio_file), possible_indexes)
+                )
+                indexes.extend(possible_indexes)
+        print("Indexes", len(indexes), "\n")
+        return indexes
+
+    def scale_and_crop(self, video):
+        h, w = video.shape[-2], video.shape[-1]
+        # scale shorter side to resolution
+
+        if self.resize_size is not None:
+            scale = self.resize_size / min(h, w)
+            if h < w:
+                target_size = (self.resize_size, math.ceil(w * scale))
+            else:
+                target_size = (math.ceil(h * scale), self.resize_size)
+            video = F.interpolate(
+                video,
+                size=target_size,
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # center crop
+            h, w = video.shape[-2], video.shape[-1]
+            w_start = (w - self.resize_size) // 2
+            h_start = (h - self.resize_size) // 2
+            video = video[
+                :,
+                :,
+                h_start : h_start + self.resize_size,
+                w_start : w_start + self.resize_size,
+            ]
+        return self.maybe_augment(video)
+
+    def _calculate_weights(self):
+        aa_processed_count = sum(
+            1
+            for item in self._indexes
+            if "AA_processed" in (item[1] if len(item) == 3 else item[0])
+        )
+        nsv_processed_count = sum(
+            1
+            for item in self._indexes
+            if "1000actors_nsv" in (item[1] if len(item) == 3 else item[0])
+        )
+        other_count = len(self._indexes) - aa_processed_count - nsv_processed_count
+
+        aa_processed_weight = 1 / aa_processed_count if aa_processed_count > 0 else 0
+        nsv_processed_weight = 1 / nsv_processed_count if nsv_processed_count > 0 else 0
+        other_weight = 1 / other_count if other_count > 0 else 0
+
+        print(
+            f"AA processed count: {aa_processed_count}, NSV processed count: {nsv_processed_count}, other count: {other_count}"
+        )
+        print(f"AA processed weight: {aa_processed_weight}")
+        print(f"NSV processed weight: {nsv_processed_weight}")
+        print(f"Other weight: {other_weight}")
+
+        weights = [
+            aa_processed_weight
+            if "AA_processed" in (item[1] if len(item) == 3 else item[0])
+            else nsv_processed_weight
+            if "1000actors_nsv" in (item[1] if len(item) == 3 else item[0])
+            else other_weight
+            for item in self._indexes
+        ]
+        return weights
+
+    def __getitem__(self, idx):
+        if self.balance_datasets:
+            idx = self.sampler.__iter__().__next__()
+
+        try:
+            (
+                original_frames,
+                clean_cond,
+                noisy_cond,
+                target,
+                audio,
+                raw_audio,
+                cond_noise,
+                emotions,
+                masks,
+                landmarks,
+            ) = self._get_frames_and_audio(idx % len(self._indexes))
+        except Exception as e:
+            print(f"Error with index {idx}: {e}")
+            return self.__getitem__(np.random.randint(0, len(self)))
+        out_data = {}
+
+        if original_frames is not None:
+            out_data["original_frames"] = original_frames
+
+        if audio is not None:
+            out_data["audio_emb"] = audio
+            out_data["raw_audio"] = raw_audio
+
+        if self.use_emotions:
+            out_data["valence"] = emotions[0]
+            out_data["arousal"] = emotions[1]
+            out_data["emo_labels"] = emotions[2]
+        if self.use_latent:
+            input_key = "latents"
+        else:
+            input_key = "frames"
+        out_data[input_key] = target
+        if noisy_cond is not None:
+            out_data["cond_frames"] = noisy_cond
+        out_data["cond_frames_without_noise"] = clean_cond
+        if cond_noise is not None:
+            out_data["cond_aug"] = cond_noise
+
+        if masks is not None:
+            out_data["masks"] = masks
+            out_data["gt"] = target
+            if landmarks is not None:
+                out_data["landmarks"] = landmarks
+
+        out_data["motion_bucket_id"] = torch.tensor([self.motion_id])
+        out_data["fps_id"] = torch.tensor([self.fps - 1])
+        out_data["num_video_frames"] = self.num_frames
+        out_data["image_only_indicator"] = torch.zeros(self.num_frames)
+        return out_data
+
+
+if __name__ == "__main__":
+    import torchvision.transforms as transforms
+    import cv2
+
+    transform = transforms.Compose(transforms=[transforms.Resize((256, 256))])
+    dataset = VideoDataset(
+        "/vol/paramonos2/projects/antoni/datasets/mahnob/filelist_videos_val.txt",
+        transform=transform,
+        num_frames=25,
+    )
+    print(len(dataset))
+    idx = np.random.randint(0, len(dataset))
+
+    for i in range(10):
+        print(dataset[i][0].shape, dataset[i][1].shape)
+
+    image_identity = (dataset[idx][0].permute(1, 2, 0).numpy() + 1) / 2 * 255
+    image_other = (dataset[idx][1][:, -1].permute(1, 2, 0).numpy() + 1) / 2 * 255
+    cv2.imwrite("image_identity.png", image_identity[:, :, ::-1])
+    for i in range(25):
+        image = (dataset[idx][1][:, i].permute(1, 2, 0).numpy() + 1) / 2 * 255
+        cv2.imwrite(f"tmp_vid_dataset/image_{i}.png", image[:, :, ::-1])

sgm/inference/api.py

@@ -0,0 +1,385 @@
+import pathlib
+from dataclasses import asdict, dataclass
+from enum import Enum
+from typing import Optional
+
+from omegaconf import OmegaConf
+
+from sgm.inference.helpers import (Img2ImgDiscretizationWrapper, do_img2img,
+                                   do_sample)
+from sgm.modules.diffusionmodules.sampling import (DPMPP2MSampler,
+                                                   DPMPP2SAncestralSampler,
+                                                   EulerAncestralSampler,
+                                                   EulerEDMSampler,
+                                                   HeunEDMSampler,
+                                                   LinearMultistepSampler)
+from sgm.util import load_model_from_config
+
+
+class ModelArchitecture(str, Enum):
+    SD_2_1 = "stable-diffusion-v2-1"
+    SD_2_1_768 = "stable-diffusion-v2-1-768"
+    SDXL_V0_9_BASE = "stable-diffusion-xl-v0-9-base"
+    SDXL_V0_9_REFINER = "stable-diffusion-xl-v0-9-refiner"
+    SDXL_V1_BASE = "stable-diffusion-xl-v1-base"
+    SDXL_V1_REFINER = "stable-diffusion-xl-v1-refiner"
+
+
+class Sampler(str, Enum):
+    EULER_EDM = "EulerEDMSampler"
+    HEUN_EDM = "HeunEDMSampler"
+    EULER_ANCESTRAL = "EulerAncestralSampler"
+    DPMPP2S_ANCESTRAL = "DPMPP2SAncestralSampler"
+    DPMPP2M = "DPMPP2MSampler"
+    LINEAR_MULTISTEP = "LinearMultistepSampler"
+
+
+class Discretization(str, Enum):
+    LEGACY_DDPM = "LegacyDDPMDiscretization"
+    EDM = "EDMDiscretization"
+
+
+class Guider(str, Enum):
+    VANILLA = "VanillaCFG"
+    IDENTITY = "IdentityGuider"
+
+
+class Thresholder(str, Enum):
+    NONE = "None"
+
+
+@dataclass
+class SamplingParams:
+    width: int = 1024
+    height: int = 1024
+    steps: int = 50
+    sampler: Sampler = Sampler.DPMPP2M
+    discretization: Discretization = Discretization.LEGACY_DDPM
+    guider: Guider = Guider.VANILLA
+    thresholder: Thresholder = Thresholder.NONE
+    scale: float = 6.0
+    aesthetic_score: float = 5.0
+    negative_aesthetic_score: float = 5.0
+    img2img_strength: float = 1.0
+    orig_width: int = 1024
+    orig_height: int = 1024
+    crop_coords_top: int = 0
+    crop_coords_left: int = 0
+    sigma_min: float = 0.0292
+    sigma_max: float = 14.6146
+    rho: float = 3.0
+    s_churn: float = 0.0
+    s_tmin: float = 0.0
+    s_tmax: float = 999.0
+    s_noise: float = 1.0
+    eta: float = 1.0
+    order: int = 4
+
+
+@dataclass
+class SamplingSpec:
+    width: int
+    height: int
+    channels: int
+    factor: int
+    is_legacy: bool
+    config: str
+    ckpt: str
+    is_guided: bool
+
+
+model_specs = {
+    ModelArchitecture.SD_2_1: SamplingSpec(
+        height=512,
+        width=512,
+        channels=4,
+        factor=8,
+        is_legacy=True,
+        config="sd_2_1.yaml",
+        ckpt="v2-1_512-ema-pruned.safetensors",
+        is_guided=True,
+    ),
+    ModelArchitecture.SD_2_1_768: SamplingSpec(
+        height=768,
+        width=768,
+        channels=4,
+        factor=8,
+        is_legacy=True,
+        config="sd_2_1_768.yaml",
+        ckpt="v2-1_768-ema-pruned.safetensors",
+        is_guided=True,
+    ),
+    ModelArchitecture.SDXL_V0_9_BASE: SamplingSpec(
+        height=1024,
+        width=1024,
+        channels=4,
+        factor=8,
+        is_legacy=False,
+        config="sd_xl_base.yaml",
+        ckpt="sd_xl_base_0.9.safetensors",
+        is_guided=True,
+    ),
+    ModelArchitecture.SDXL_V0_9_REFINER: SamplingSpec(
+        height=1024,
+        width=1024,
+        channels=4,
+        factor=8,
+        is_legacy=True,
+        config="sd_xl_refiner.yaml",
+        ckpt="sd_xl_refiner_0.9.safetensors",
+        is_guided=True,
+    ),
+    ModelArchitecture.SDXL_V1_BASE: SamplingSpec(
+        height=1024,
+        width=1024,
+        channels=4,
+        factor=8,
+        is_legacy=False,
+        config="sd_xl_base.yaml",
+        ckpt="sd_xl_base_1.0.safetensors",
+        is_guided=True,
+    ),
+    ModelArchitecture.SDXL_V1_REFINER: SamplingSpec(
+        height=1024,
+        width=1024,
+        channels=4,
+        factor=8,
+        is_legacy=True,
+        config="sd_xl_refiner.yaml",
+        ckpt="sd_xl_refiner_1.0.safetensors",
+        is_guided=True,
+    ),
+}
+
+
+class SamplingPipeline:
+    def __init__(
+        self,
+        model_id: ModelArchitecture,
+        model_path="checkpoints",
+        config_path="configs/inference",
+        device="cuda",
+        use_fp16=True,
+    ) -> None:
+        if model_id not in model_specs:
+            raise ValueError(f"Model {model_id} not supported")
+        self.model_id = model_id
+        self.specs = model_specs[self.model_id]
+        self.config = str(pathlib.Path(config_path, self.specs.config))
+        self.ckpt = str(pathlib.Path(model_path, self.specs.ckpt))
+        self.device = device
+        self.model = self._load_model(device=device, use_fp16=use_fp16)
+
+    def _load_model(self, device="cuda", use_fp16=True):
+        config = OmegaConf.load(self.config)
+        model = load_model_from_config(config, self.ckpt)
+        if model is None:
+            raise ValueError(f"Model {self.model_id} could not be loaded")
+        model.to(device)
+        if use_fp16:
+            model.conditioner.half()
+            model.model.half()
+        return model
+
+    def text_to_image(
+        self,
+        params: SamplingParams,
+        prompt: str,
+        negative_prompt: str = "",
+        samples: int = 1,
+        return_latents: bool = False,
+    ):
+        sampler = get_sampler_config(params)
+        value_dict = asdict(params)
+        value_dict["prompt"] = prompt
+        value_dict["negative_prompt"] = negative_prompt
+        value_dict["target_width"] = params.width
+        value_dict["target_height"] = params.height
+        return do_sample(
+            self.model,
+            sampler,
+            value_dict,
+            samples,
+            params.height,
+            params.width,
+            self.specs.channels,
+            self.specs.factor,
+            force_uc_zero_embeddings=["txt"] if not self.specs.is_legacy else [],
+            return_latents=return_latents,
+            filter=None,
+        )
+
+    def image_to_image(
+        self,
+        params: SamplingParams,
+        image,
+        prompt: str,
+        negative_prompt: str = "",
+        samples: int = 1,
+        return_latents: bool = False,
+    ):
+        sampler = get_sampler_config(params)
+
+        if params.img2img_strength < 1.0:
+            sampler.discretization = Img2ImgDiscretizationWrapper(
+                sampler.discretization,
+                strength=params.img2img_strength,
+            )
+        height, width = image.shape[2], image.shape[3]
+        value_dict = asdict(params)
+        value_dict["prompt"] = prompt
+        value_dict["negative_prompt"] = negative_prompt
+        value_dict["target_width"] = width
+        value_dict["target_height"] = height
+        return do_img2img(
+            image,
+            self.model,
+            sampler,
+            value_dict,
+            samples,
+            force_uc_zero_embeddings=["txt"] if not self.specs.is_legacy else [],
+            return_latents=return_latents,
+            filter=None,
+        )
+
+    def refiner(
+        self,
+        params: SamplingParams,
+        image,
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        samples: int = 1,
+        return_latents: bool = False,
+    ):
+        sampler = get_sampler_config(params)
+        value_dict = {
+            "orig_width": image.shape[3] * 8,
+            "orig_height": image.shape[2] * 8,
+            "target_width": image.shape[3] * 8,
+            "target_height": image.shape[2] * 8,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "crop_coords_top": 0,
+            "crop_coords_left": 0,
+            "aesthetic_score": 6.0,
+            "negative_aesthetic_score": 2.5,
+        }
+
+        return do_img2img(
+            image,
+            self.model,
+            sampler,
+            value_dict,
+            samples,
+            skip_encode=True,
+            return_latents=return_latents,
+            filter=None,
+        )
+
+
+def get_guider_config(params: SamplingParams):
+    if params.guider == Guider.IDENTITY:
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.IdentityGuider"
+        }
+    elif params.guider == Guider.VANILLA:
+        scale = params.scale
+
+        thresholder = params.thresholder
+
+        if thresholder == Thresholder.NONE:
+            dyn_thresh_config = {
+                "target": "sgm.modules.diffusionmodules.sampling_utils.NoDynamicThresholding"
+            }
+        else:
+            raise NotImplementedError
+
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.VanillaCFG",
+            "params": {"scale": scale, "dyn_thresh_config": dyn_thresh_config},
+        }
+    else:
+        raise NotImplementedError
+    return guider_config
+
+
+def get_discretization_config(params: SamplingParams):
+    if params.discretization == Discretization.LEGACY_DDPM:
+        discretization_config = {
+            "target": "sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization",
+        }
+    elif params.discretization == Discretization.EDM:
+        discretization_config = {
+            "target": "sgm.modules.diffusionmodules.discretizer.EDMDiscretization",
+            "params": {
+                "sigma_min": params.sigma_min,
+                "sigma_max": params.sigma_max,
+                "rho": params.rho,
+            },
+        }
+    else:
+        raise ValueError(f"unknown discretization {params.discretization}")
+    return discretization_config
+
+
+def get_sampler_config(params: SamplingParams):
+    discretization_config = get_discretization_config(params)
+    guider_config = get_guider_config(params)
+    sampler = None
+    if params.sampler == Sampler.EULER_EDM:
+        return EulerEDMSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            s_churn=params.s_churn,
+            s_tmin=params.s_tmin,
+            s_tmax=params.s_tmax,
+            s_noise=params.s_noise,
+            verbose=True,
+        )
+    if params.sampler == Sampler.HEUN_EDM:
+        return HeunEDMSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            s_churn=params.s_churn,
+            s_tmin=params.s_tmin,
+            s_tmax=params.s_tmax,
+            s_noise=params.s_noise,
+            verbose=True,
+        )
+    if params.sampler == Sampler.EULER_ANCESTRAL:
+        return EulerAncestralSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            eta=params.eta,
+            s_noise=params.s_noise,
+            verbose=True,
+        )
+    if params.sampler == Sampler.DPMPP2S_ANCESTRAL:
+        return DPMPP2SAncestralSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            eta=params.eta,
+            s_noise=params.s_noise,
+            verbose=True,
+        )
+    if params.sampler == Sampler.DPMPP2M:
+        return DPMPP2MSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            verbose=True,
+        )
+    if params.sampler == Sampler.LINEAR_MULTISTEP:
+        return LinearMultistepSampler(
+            num_steps=params.steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            order=params.order,
+            verbose=True,
+        )
+
+    raise ValueError(f"unknown sampler {params.sampler}!")

sgm/inference/helpers.py

@@ -0,0 +1,305 @@
+import math
+import os
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from einops import rearrange
+from imwatermark import WatermarkEncoder
+from omegaconf import ListConfig
+from PIL import Image
+from torch import autocast
+
+from sgm.util import append_dims
+
+
+class WatermarkEmbedder:
+    def __init__(self, watermark):
+        self.watermark = watermark
+        self.num_bits = len(WATERMARK_BITS)
+        self.encoder = WatermarkEncoder()
+        self.encoder.set_watermark("bits", self.watermark)
+
+    def __call__(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Adds a predefined watermark to the input image
+
+        Args:
+            image: ([N,] B, RGB, H, W) in range [0, 1]
+
+        Returns:
+            same as input but watermarked
+        """
+        squeeze = len(image.shape) == 4
+        if squeeze:
+            image = image[None, ...]
+        n = image.shape[0]
+        image_np = rearrange(
+            (255 * image).detach().cpu(), "n b c h w -> (n b) h w c"
+        ).numpy()[:, :, :, ::-1]
+        # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
+        # watermarking libary expects input as cv2 BGR format
+        for k in range(image_np.shape[0]):
+            image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
+        image = torch.from_numpy(
+            rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)
+        ).to(image.device)
+        image = torch.clamp(image / 255, min=0.0, max=1.0)
+        if squeeze:
+            image = image[0]
+        return image
+
+
+# A fixed 48-bit message that was choosen at random
+# WATERMARK_MESSAGE = 0xB3EC907BB19E
+WATERMARK_MESSAGE = 0b101100111110110010010000011110111011000110011110
+# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
+WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
+embed_watermark = WatermarkEmbedder(WATERMARK_BITS)
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    return list({x.input_key for x in conditioner.embedders})
+
+
+def perform_save_locally(save_path, samples):
+    os.makedirs(os.path.join(save_path), exist_ok=True)
+    base_count = len(os.listdir(os.path.join(save_path)))
+    samples = embed_watermark(samples)
+    for sample in samples:
+        sample = 255.0 * rearrange(sample.cpu().numpy(), "c h w -> h w c")
+        Image.fromarray(sample.astype(np.uint8)).save(
+            os.path.join(save_path, f"{base_count:09}.png")
+        )
+        base_count += 1
+
+
+class Img2ImgDiscretizationWrapper:
+    """
+    wraps a discretizer, and prunes the sigmas
+    params:
+        strength: float between 0.0 and 1.0. 1.0 means full sampling (all sigmas are returned)
+    """
+
+    def __init__(self, discretization, strength: float = 1.0):
+        self.discretization = discretization
+        self.strength = strength
+        assert 0.0 <= self.strength <= 1.0
+
+    def __call__(self, *args, **kwargs):
+        # sigmas start large first, and decrease then
+        sigmas = self.discretization(*args, **kwargs)
+        print(f"sigmas after discretization, before pruning img2img: ", sigmas)
+        sigmas = torch.flip(sigmas, (0,))
+        sigmas = sigmas[: max(int(self.strength * len(sigmas)), 1)]
+        print("prune index:", max(int(self.strength * len(sigmas)), 1))
+        sigmas = torch.flip(sigmas, (0,))
+        print(f"sigmas after pruning: ", sigmas)
+        return sigmas
+
+
+def do_sample(
+    model,
+    sampler,
+    value_dict,
+    num_samples,
+    H,
+    W,
+    C,
+    F,
+    force_uc_zero_embeddings: Optional[List] = None,
+    batch2model_input: Optional[List] = None,
+    return_latents=False,
+    filter=None,
+    device="cuda",
+):
+    if force_uc_zero_embeddings is None:
+        force_uc_zero_embeddings = []
+    if batch2model_input is None:
+        batch2model_input = []
+
+    with torch.no_grad():
+        with autocast(device) as precision_scope:
+            with model.ema_scope():
+                num_samples = [num_samples]
+                batch, batch_uc = get_batch(
+                    get_unique_embedder_keys_from_conditioner(model.conditioner),
+                    value_dict,
+                    num_samples,
+                )
+                for key in batch:
+                    if isinstance(batch[key], torch.Tensor):
+                        print(key, batch[key].shape)
+                    elif isinstance(batch[key], list):
+                        print(key, [len(l) for l in batch[key]])
+                    else:
+                        print(key, batch[key])
+                c, uc = model.conditioner.get_unconditional_conditioning(
+                    batch,
+                    batch_uc=batch_uc,
+                    force_uc_zero_embeddings=force_uc_zero_embeddings,
+                )
+
+                for k in c:
+                    if not k == "crossattn":
+                        c[k], uc[k] = map(
+                            lambda y: y[k][: math.prod(num_samples)].to(device), (c, uc)
+                        )
+
+                additional_model_inputs = {}
+                for k in batch2model_input:
+                    additional_model_inputs[k] = batch[k]
+
+                shape = (math.prod(num_samples), C, H // F, W // F)
+                randn = torch.randn(shape).to(device)
+
+                def denoiser(input, sigma, c):
+                    return model.denoiser(
+                        model.model, input, sigma, c, **additional_model_inputs
+                    )
+
+                samples_z = sampler(denoiser, randn, cond=c, uc=uc)
+                samples_x = model.decode_first_stage(samples_z)
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+
+                if filter is not None:
+                    samples = filter(samples)
+
+                if return_latents:
+                    return samples, samples_z
+                return samples
+
+
+def get_batch(keys, value_dict, N: Union[List, ListConfig], device="cuda"):
+    # Hardcoded demo setups; might undergo some changes in the future
+
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "txt":
+            batch["txt"] = (
+                np.repeat([value_dict["prompt"]], repeats=math.prod(N))
+                .reshape(N)
+                .tolist()
+            )
+            batch_uc["txt"] = (
+                np.repeat([value_dict["negative_prompt"]], repeats=math.prod(N))
+                .reshape(N)
+                .tolist()
+            )
+        elif key == "original_size_as_tuple":
+            batch["original_size_as_tuple"] = (
+                torch.tensor([value_dict["orig_height"], value_dict["orig_width"]])
+                .to(device)
+                .repeat(*N, 1)
+            )
+        elif key == "crop_coords_top_left":
+            batch["crop_coords_top_left"] = (
+                torch.tensor(
+                    [value_dict["crop_coords_top"], value_dict["crop_coords_left"]]
+                )
+                .to(device)
+                .repeat(*N, 1)
+            )
+        elif key == "aesthetic_score":
+            batch["aesthetic_score"] = (
+                torch.tensor([value_dict["aesthetic_score"]]).to(device).repeat(*N, 1)
+            )
+            batch_uc["aesthetic_score"] = (
+                torch.tensor([value_dict["negative_aesthetic_score"]])
+                .to(device)
+                .repeat(*N, 1)
+            )
+
+        elif key == "target_size_as_tuple":
+            batch["target_size_as_tuple"] = (
+                torch.tensor([value_dict["target_height"], value_dict["target_width"]])
+                .to(device)
+                .repeat(*N, 1)
+            )
+        else:
+            batch[key] = value_dict[key]
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+    return batch, batch_uc
+
+
+def get_input_image_tensor(image: Image.Image, device="cuda"):
+    w, h = image.size
+    print(f"loaded input image of size ({w}, {h})")
+    width, height = map(
+        lambda x: x - x % 64, (w, h)
+    )  # resize to integer multiple of 64
+    image = image.resize((width, height))
+    image_array = np.array(image.convert("RGB"))
+    image_array = image_array[None].transpose(0, 3, 1, 2)
+    image_tensor = torch.from_numpy(image_array).to(dtype=torch.float32) / 127.5 - 1.0
+    return image_tensor.to(device)
+
+
+def do_img2img(
+    img,
+    model,
+    sampler,
+    value_dict,
+    num_samples,
+    force_uc_zero_embeddings=[],
+    additional_kwargs={},
+    offset_noise_level: float = 0.0,
+    return_latents=False,
+    skip_encode=False,
+    filter=None,
+    device="cuda",
+):
+    with torch.no_grad():
+        with autocast(device) as precision_scope:
+            with model.ema_scope():
+                batch, batch_uc = get_batch(
+                    get_unique_embedder_keys_from_conditioner(model.conditioner),
+                    value_dict,
+                    [num_samples],
+                )
+                c, uc = model.conditioner.get_unconditional_conditioning(
+                    batch,
+                    batch_uc=batch_uc,
+                    force_uc_zero_embeddings=force_uc_zero_embeddings,
+                )
+
+                for k in c:
+                    c[k], uc[k] = map(lambda y: y[k][:num_samples].to(device), (c, uc))
+
+                for k in additional_kwargs:
+                    c[k] = uc[k] = additional_kwargs[k]
+                if skip_encode:
+                    z = img
+                else:
+                    z = model.encode_first_stage(img)
+                noise = torch.randn_like(z)
+                sigmas = sampler.discretization(sampler.num_steps)
+                sigma = sigmas[0].to(z.device)
+
+                if offset_noise_level > 0.0:
+                    noise = noise + offset_noise_level * append_dims(
+                        torch.randn(z.shape[0], device=z.device), z.ndim
+                    )
+                noised_z = z + noise * append_dims(sigma, z.ndim)
+                noised_z = noised_z / torch.sqrt(
+                    1.0 + sigmas[0] ** 2.0
+                )  # Note: hardcoded to DDPM-like scaling. need to generalize later.
+
+                def denoiser(x, sigma, c):
+                    return model.denoiser(model.model, x, sigma, c)
+
+                samples_z = sampler(denoiser, noised_z, cond=c, uc=uc)
+                samples_x = model.decode_first_stage(samples_z)
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+
+                if filter is not None:
+                    samples = filter(samples)
+
+                if return_latents:
+                    return samples, samples_z
+                return samples

sgm/lr_scheduler.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+
+    def __init__(
+        self,
+        warm_up_steps,
+        lr_min,
+        lr_max,
+        lr_start,
+        max_decay_steps,
+        verbosity_interval=0,
+    ):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (
+                self.lr_max - self.lr_start
+            ) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (
+                self.lr_max_decay_steps - self.lr_warm_up_steps
+            )
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                1 + np.cos(t * np.pi)
+            )
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+
+    def __init__(
+        self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0
+    ):
+        assert (
+            len(warm_up_steps)
+            == len(f_min)
+            == len(f_max)
+            == len(f_start)
+            == len(cycle_lengths)
+        )
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(
+                    f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                    f"current cycle {cycle}"
+                )
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+                cycle
+            ] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (
+                self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]
+            )
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                1 + np.cos(t * np.pi)
+            )
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(
+                    f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                    f"current cycle {cycle}"
+                )
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+                cycle
+            ] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (
+                self.cycle_lengths[cycle] - n
+            ) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f

sgm/models/__init__.py

@@ -0,0 +1,2 @@
+from .autoencoder import AutoencodingEngine
+from .diffusion import DiffusionEngine

sgm/models/autoencoder.py

@@ -0,0 +1,615 @@
+import logging
+import math
+import re
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from einops import rearrange
+from packaging import version
+
+from ..modules.autoencoding.regularizers import AbstractRegularizer
+from ..modules.ema import LitEma
+from ..util import (default, get_nested_attribute, get_obj_from_str,
+                    instantiate_from_config)
+
+logpy = logging.getLogger(__name__)
+
+
+class AbstractAutoencoder(pl.LightningModule):
+    """
+    This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators,
+    unCLIP models, etc. Hence, it is fairly general, and specific features
+    (e.g. discriminator training, encoding, decoding) must be implemented in subclasses.
+    """
+
+    def __init__(
+        self,
+        ema_decay: Union[None, float] = None,
+        monitor: Union[None, str] = None,
+        input_key: str = "jpg",
+    ):
+        super().__init__()
+
+        self.input_key = input_key
+        self.use_ema = ema_decay is not None
+        if monitor is not None:
+            self.monitor = monitor
+
+        if self.use_ema:
+            self.model_ema = LitEma(self, decay=ema_decay)
+            logpy.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            self.automatic_optimization = False
+
+    def apply_ckpt(self, ckpt: Union[None, str, dict]):
+        if ckpt is None:
+            return
+        if isinstance(ckpt, str):
+            ckpt = {
+                "target": "sgm.modules.checkpoint.CheckpointEngine",
+                "params": {"ckpt_path": ckpt},
+            }
+        engine = instantiate_from_config(ckpt)
+        engine(self)
+
+    @abstractmethod
+    def get_input(self, batch) -> Any:
+        raise NotImplementedError()
+
+    def on_train_batch_end(self, *args, **kwargs):
+        # for EMA computation
+        if self.use_ema:
+            self.model_ema(self)
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.parameters())
+            self.model_ema.copy_to(self)
+            if context is not None:
+                logpy.info(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.parameters())
+                if context is not None:
+                    logpy.info(f"{context}: Restored training weights")
+
+    @abstractmethod
+    def encode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("encode()-method of abstract base class called")
+
+    @abstractmethod
+    def decode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("decode()-method of abstract base class called")
+
+    def instantiate_optimizer_from_config(self, params, lr, cfg):
+        logpy.info(f"loading >>> {cfg['target']} <<< optimizer from config")
+        return get_obj_from_str(cfg["target"])(
+            params, lr=lr, **cfg.get("params", dict())
+        )
+
+    def configure_optimizers(self) -> Any:
+        raise NotImplementedError()
+
+
+class AutoencodingEngine(AbstractAutoencoder):
+    """
+    Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL
+    (we also restore them explicitly as special cases for legacy reasons).
+    Regularizations such as KL or VQ are moved to the regularizer class.
+    """
+
+    def __init__(
+        self,
+        *args,
+        encoder_config: Dict,
+        decoder_config: Dict,
+        loss_config: Dict,
+        regularizer_config: Dict,
+        optimizer_config: Union[Dict, None] = None,
+        lr_g_factor: float = 1.0,
+        trainable_ae_params: Optional[List[List[str]]] = None,
+        ae_optimizer_args: Optional[List[dict]] = None,
+        trainable_disc_params: Optional[List[List[str]]] = None,
+        disc_optimizer_args: Optional[List[dict]] = None,
+        disc_start_iter: int = 0,
+        diff_boost_factor: float = 3.0,
+        ckpt_engine: Union[None, str, dict] = None,
+        ckpt_path: Optional[str] = None,
+        additional_decode_keys: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.automatic_optimization = False  # pytorch lightning
+
+        self.encoder: torch.nn.Module = instantiate_from_config(encoder_config)
+        self.decoder: torch.nn.Module = instantiate_from_config(decoder_config)
+        self.loss: torch.nn.Module = instantiate_from_config(loss_config)
+        self.regularization: AbstractRegularizer = instantiate_from_config(
+            regularizer_config
+        )
+        self.optimizer_config = default(
+            optimizer_config, {"target": "torch.optim.Adam"}
+        )
+        self.diff_boost_factor = diff_boost_factor
+        self.disc_start_iter = disc_start_iter
+        self.lr_g_factor = lr_g_factor
+        self.trainable_ae_params = trainable_ae_params
+        if self.trainable_ae_params is not None:
+            self.ae_optimizer_args = default(
+                ae_optimizer_args,
+                [{} for _ in range(len(self.trainable_ae_params))],
+            )
+            assert len(self.ae_optimizer_args) == len(self.trainable_ae_params)
+        else:
+            self.ae_optimizer_args = [{}]  # makes type consitent
+
+        self.trainable_disc_params = trainable_disc_params
+        if self.trainable_disc_params is not None:
+            self.disc_optimizer_args = default(
+                disc_optimizer_args,
+                [{} for _ in range(len(self.trainable_disc_params))],
+            )
+            assert len(self.disc_optimizer_args) == len(self.trainable_disc_params)
+        else:
+            self.disc_optimizer_args = [{}]  # makes type consitent
+
+        if ckpt_path is not None:
+            assert ckpt_engine is None, "Can't set ckpt_engine and ckpt_path"
+            logpy.warn("Checkpoint path is deprecated, use `checkpoint_egnine` instead")
+        self.apply_ckpt(default(ckpt_path, ckpt_engine))
+        self.additional_decode_keys = set(default(additional_decode_keys, []))
+
+    def get_input(self, batch: Dict) -> torch.Tensor:
+        # assuming unified data format, dataloader returns a dict.
+        # image tensors should be scaled to -1 ... 1 and in channels-first
+        # format (e.g., bchw instead if bhwc)
+        return batch[self.input_key]
+
+    def get_autoencoder_params(self) -> list:
+        params = []
+        if hasattr(self.loss, "get_trainable_autoencoder_parameters"):
+            params += list(self.loss.get_trainable_autoencoder_parameters())
+        if hasattr(self.regularization, "get_trainable_parameters"):
+            params += list(self.regularization.get_trainable_parameters())
+        params = params + list(self.encoder.parameters())
+        params = params + list(self.decoder.parameters())
+        return params
+
+    def get_discriminator_params(self) -> list:
+        if hasattr(self.loss, "get_trainable_parameters"):
+            params = list(self.loss.get_trainable_parameters())  # e.g., discriminator
+        else:
+            params = []
+        return params
+
+    def get_last_layer(self):
+        return self.decoder.get_last_layer()
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        return_reg_log: bool = False,
+        unregularized: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        z = self.encoder(x)
+        if unregularized:
+            return z, dict()
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(self, z: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.decoder(z, **kwargs)
+        return x
+
+    def forward(
+        self, x: torch.Tensor, **additional_decode_kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor, dict]:
+        z, reg_log = self.encode(x, return_reg_log=True)
+        dec = self.decode(z, **additional_decode_kwargs)
+        return z, dec, reg_log
+
+    def inner_training_step(
+        self, batch: dict, batch_idx: int, optimizer_idx: int = 0
+    ) -> torch.Tensor:
+        x = self.get_input(batch)
+        additional_decode_kwargs = {
+            key: batch[key] for key in self.additional_decode_keys.intersection(batch)
+        }
+        z, xrec, regularization_log = self(x, **additional_decode_kwargs)
+        if hasattr(self.loss, "forward_keys"):
+            extra_info = {
+                "z": z,
+                "optimizer_idx": optimizer_idx,
+                "global_step": self.global_step,
+                "last_layer": self.get_last_layer(),
+                "split": "train",
+                "regularization_log": regularization_log,
+                "autoencoder": self,
+            }
+            extra_info = {k: extra_info[k] for k in self.loss.forward_keys}
+        else:
+            extra_info = dict()
+
+        if optimizer_idx == 0:
+            # autoencode
+            out_loss = self.loss(x, xrec, **extra_info)
+            if isinstance(out_loss, tuple):
+                aeloss, log_dict_ae = out_loss
+            else:
+                # simple loss function
+                aeloss = out_loss
+                log_dict_ae = {"train/loss/rec": aeloss.detach()}
+
+            self.log_dict(
+                log_dict_ae,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=False,
+            )
+            self.log(
+                "loss",
+                aeloss.mean().detach(),
+                prog_bar=True,
+                logger=False,
+                on_epoch=False,
+                on_step=True,
+            )
+            return aeloss
+        elif optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(x, xrec, **extra_info)
+            # -> discriminator always needs to return a tuple
+            self.log_dict(
+                log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True
+            )
+            return discloss
+        else:
+            raise NotImplementedError(f"Unknown optimizer {optimizer_idx}")
+
+    def training_step(self, batch: dict, batch_idx: int):
+        opts = self.optimizers()
+        if not isinstance(opts, list):
+            # Non-adversarial case
+            opts = [opts]
+        optimizer_idx = batch_idx % len(opts)
+        if self.global_step < self.disc_start_iter:
+            optimizer_idx = 0
+        opt = opts[optimizer_idx]
+        opt.zero_grad()
+        with opt.toggle_model():
+            loss = self.inner_training_step(
+                batch, batch_idx, optimizer_idx=optimizer_idx
+            )
+            self.manual_backward(loss)
+        opt.step()
+
+    def validation_step(self, batch: dict, batch_idx: int) -> Dict:
+        log_dict = self._validation_step(batch, batch_idx)
+        with self.ema_scope():
+            log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
+            log_dict.update(log_dict_ema)
+        return log_dict
+
+    def _validation_step(self, batch: dict, batch_idx: int, postfix: str = "") -> Dict:
+        x = self.get_input(batch)
+
+        z, xrec, regularization_log = self(x)
+        if hasattr(self.loss, "forward_keys"):
+            extra_info = {
+                "z": z,
+                "optimizer_idx": 0,
+                "global_step": self.global_step,
+                "last_layer": self.get_last_layer(),
+                "split": "val" + postfix,
+                "regularization_log": regularization_log,
+                "autoencoder": self,
+            }
+            extra_info = {k: extra_info[k] for k in self.loss.forward_keys}
+        else:
+            extra_info = dict()
+        out_loss = self.loss(x, xrec, **extra_info)
+        if isinstance(out_loss, tuple):
+            aeloss, log_dict_ae = out_loss
+        else:
+            # simple loss function
+            aeloss = out_loss
+            log_dict_ae = {f"val{postfix}/loss/rec": aeloss.detach()}
+        full_log_dict = log_dict_ae
+
+        if "optimizer_idx" in extra_info:
+            extra_info["optimizer_idx"] = 1
+            discloss, log_dict_disc = self.loss(x, xrec, **extra_info)
+            full_log_dict.update(log_dict_disc)
+        self.log(
+            f"val{postfix}/loss/rec",
+            log_dict_ae[f"val{postfix}/loss/rec"],
+            sync_dist=True,
+        )
+        self.log_dict(full_log_dict, sync_dist=True)
+        return full_log_dict
+
+    def get_param_groups(
+        self, parameter_names: List[List[str]], optimizer_args: List[dict]
+    ) -> Tuple[List[Dict[str, Any]], int]:
+        groups = []
+        num_params = 0
+        for names, args in zip(parameter_names, optimizer_args):
+            params = []
+            for pattern_ in names:
+                pattern_params = []
+                pattern = re.compile(pattern_)
+                for p_name, param in self.named_parameters():
+                    if re.match(pattern, p_name):
+                        pattern_params.append(param)
+                        num_params += param.numel()
+                if len(pattern_params) == 0:
+                    logpy.warn(f"Did not find parameters for pattern {pattern_}")
+                params.extend(pattern_params)
+            groups.append({"params": params, **args})
+        return groups, num_params
+
+    def configure_optimizers(self) -> List[torch.optim.Optimizer]:
+        if self.trainable_ae_params is None:
+            ae_params = self.get_autoencoder_params()
+        else:
+            ae_params, num_ae_params = self.get_param_groups(
+                self.trainable_ae_params, self.ae_optimizer_args
+            )
+            logpy.info(f"Number of trainable autoencoder parameters: {num_ae_params:,}")
+        if self.trainable_disc_params is None:
+            disc_params = self.get_discriminator_params()
+        else:
+            disc_params, num_disc_params = self.get_param_groups(
+                self.trainable_disc_params, self.disc_optimizer_args
+            )
+            logpy.info(
+                f"Number of trainable discriminator parameters: {num_disc_params:,}"
+            )
+        opt_ae = self.instantiate_optimizer_from_config(
+            ae_params,
+            default(self.lr_g_factor, 1.0) * self.learning_rate,
+            self.optimizer_config,
+        )
+        opts = [opt_ae]
+        if len(disc_params) > 0:
+            opt_disc = self.instantiate_optimizer_from_config(
+                disc_params, self.learning_rate, self.optimizer_config
+            )
+            opts.append(opt_disc)
+
+        return opts
+
+    @torch.no_grad()
+    def log_images(
+        self, batch: dict, additional_log_kwargs: Optional[Dict] = None, **kwargs
+    ) -> dict:
+        log = dict()
+        additional_decode_kwargs = {}
+        x = self.get_input(batch)
+        additional_decode_kwargs.update(
+            {key: batch[key] for key in self.additional_decode_keys.intersection(batch)}
+        )
+
+        _, xrec, _ = self(x, **additional_decode_kwargs)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        diff = 0.5 * torch.abs(torch.clamp(xrec, -1.0, 1.0) - x)
+        diff.clamp_(0, 1.0)
+        log["diff"] = 2.0 * diff - 1.0
+        # diff_boost shows location of small errors, by boosting their
+        # brightness.
+        log["diff_boost"] = (
+            2.0 * torch.clamp(self.diff_boost_factor * diff, 0.0, 1.0) - 1
+        )
+        if hasattr(self.loss, "log_images"):
+            log.update(self.loss.log_images(x, xrec))
+        with self.ema_scope():
+            _, xrec_ema, _ = self(x, **additional_decode_kwargs)
+            log["reconstructions_ema"] = xrec_ema
+            diff_ema = 0.5 * torch.abs(torch.clamp(xrec_ema, -1.0, 1.0) - x)
+            diff_ema.clamp_(0, 1.0)
+            log["diff_ema"] = 2.0 * diff_ema - 1.0
+            log["diff_boost_ema"] = (
+                2.0 * torch.clamp(self.diff_boost_factor * diff_ema, 0.0, 1.0) - 1
+            )
+        if additional_log_kwargs:
+            additional_decode_kwargs.update(additional_log_kwargs)
+            _, xrec_add, _ = self(x, **additional_decode_kwargs)
+            log_str = "reconstructions-" + "-".join(
+                [f"{key}={additional_log_kwargs[key]}" for key in additional_log_kwargs]
+            )
+            log[log_str] = xrec_add
+        return log
+
+
+class AutoencodingEngineLegacy(AutoencodingEngine):
+    def __init__(self, embed_dim: int, **kwargs):
+        self.max_batch_size = kwargs.pop("max_batch_size", None)
+        ddconfig = kwargs.pop("ddconfig")
+        ckpt_path = kwargs.pop("ckpt_path", None)
+        ckpt_engine = kwargs.pop("ckpt_engine", None)
+        super().__init__(
+            encoder_config={
+                "target": "sgm.modules.diffusionmodules.model.Encoder",
+                "params": ddconfig,
+            },
+            decoder_config={
+                "target": "sgm.modules.diffusionmodules.model.Decoder",
+                "params": ddconfig,
+            },
+            **kwargs,
+        )
+        self.quant_conv = torch.nn.Conv2d(
+            (1 + ddconfig["double_z"]) * ddconfig["z_channels"],
+            (1 + ddconfig["double_z"]) * embed_dim,
+            1,
+        )
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+
+        self.apply_ckpt(default(ckpt_path, ckpt_engine))
+
+    def get_autoencoder_params(self) -> list:
+        params = super().get_autoencoder_params()
+        return params
+
+    def encode(
+        self, x: torch.Tensor, return_reg_log: bool = False
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        if self.max_batch_size is None:
+            z = self.encoder(x)
+            z = self.quant_conv(z)
+        else:
+            N = x.shape[0]
+            bs = self.max_batch_size
+            n_batches = int(math.ceil(N / bs))
+            z = list()
+            for i_batch in range(n_batches):
+                z_batch = self.encoder(x[i_batch * bs : (i_batch + 1) * bs])
+                z_batch = self.quant_conv(z_batch)
+                z.append(z_batch)
+            z = torch.cat(z, 0)
+
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
+        if self.max_batch_size is None:
+            dec = self.post_quant_conv(z)
+            dec = self.decoder(dec, **decoder_kwargs)
+        else:
+            N = z.shape[0]
+            bs = self.max_batch_size
+            n_batches = int(math.ceil(N / bs))
+            dec = list()
+            for i_batch in range(n_batches):
+                dec_batch = self.post_quant_conv(z[i_batch * bs : (i_batch + 1) * bs])
+                dec_batch = self.decoder(dec_batch, **decoder_kwargs)
+                dec.append(dec_batch)
+            dec = torch.cat(dec, 0)
+
+        return dec
+
+
+class AutoencoderKL(AutoencodingEngineLegacy):
+    def __init__(self, **kwargs):
+        if "lossconfig" in kwargs:
+            kwargs["loss_config"] = kwargs.pop("lossconfig")
+        super().__init__(
+            regularizer_config={
+                "target": (
+                    "sgm.modules.autoencoding.regularizers"
+                    ".DiagonalGaussianRegularizer"
+                )
+            },
+            **kwargs,
+        )
+
+
+class AutoencoderLegacyVQ(AutoencodingEngineLegacy):
+    def __init__(
+        self,
+        embed_dim: int,
+        n_embed: int,
+        sane_index_shape: bool = False,
+        **kwargs,
+    ):
+        if "lossconfig" in kwargs:
+            logpy.warn(f"Parameter `lossconfig` is deprecated, use `loss_config`.")
+            kwargs["loss_config"] = kwargs.pop("lossconfig")
+        super().__init__(
+            regularizer_config={
+                "target": (
+                    "sgm.modules.autoencoding.regularizers.quantize" ".VectorQuantizer"
+                ),
+                "params": {
+                    "n_e": n_embed,
+                    "e_dim": embed_dim,
+                    "sane_index_shape": sane_index_shape,
+                },
+            },
+            **kwargs,
+        )
+
+
+class IdentityFirstStage(AbstractAutoencoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_input(self, x: Any) -> Any:
+        return x
+
+    def encode(self, x: Any, *args, **kwargs) -> Any:
+        return x
+
+    def decode(self, x: Any, *args, **kwargs) -> Any:
+        return x
+
+
+class AEIntegerWrapper(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        shape: Union[None, Tuple[int, int], List[int]] = (16, 16),
+        regularization_key: str = "regularization",
+        encoder_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        self.model = model
+        assert hasattr(model, "encode") and hasattr(
+            model, "decode"
+        ), "Need AE interface"
+        self.regularization = get_nested_attribute(model, regularization_key)
+        self.shape = shape
+        self.encoder_kwargs = default(encoder_kwargs, {"return_reg_log": True})
+
+    def encode(self, x) -> torch.Tensor:
+        assert (
+            not self.training
+        ), f"{self.__class__.__name__} only supports inference currently"
+        _, log = self.model.encode(x, **self.encoder_kwargs)
+        assert isinstance(log, dict)
+        inds = log["min_encoding_indices"]
+        return rearrange(inds, "b ... -> b (...)")
+
+    def decode(
+        self, inds: torch.Tensor, shape: Union[None, tuple, list] = None
+    ) -> torch.Tensor:
+        # expect inds shape (b, s) with s = h*w
+        shape = default(shape, self.shape)  # Optional[(h, w)]
+        if shape is not None:
+            assert len(shape) == 2, f"Unhandeled shape {shape}"
+            inds = rearrange(inds, "b (h w) -> b h w", h=shape[0], w=shape[1])
+        h = self.regularization.get_codebook_entry(inds)  # (b, h, w, c)
+        h = rearrange(h, "b h w c -> b c h w")
+        return self.model.decode(h)
+
+
+class AutoencoderKLModeOnly(AutoencodingEngineLegacy):
+    def __init__(self, **kwargs):
+        if "lossconfig" in kwargs:
+            kwargs["loss_config"] = kwargs.pop("lossconfig")
+        super().__init__(
+            regularizer_config={
+                "target": (
+                    "sgm.modules.autoencoding.regularizers"
+                    ".DiagonalGaussianRegularizer"
+                ),
+                "params": {"sample": False},
+            },
+            **kwargs,
+        )

sgm/models/diffusion.py

@@ -0,0 +1,747 @@
+import os
+import math
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Tuple, Union
+import re
+import pytorch_lightning as pl
+import torch
+from omegaconf import ListConfig, OmegaConf
+from safetensors.torch import load_file as load_safetensors
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange
+from diffusers.models.attention_processor import IPAdapterAttnProcessor2_0
+
+from ..modules import UNCONDITIONAL_CONFIG
+from ..modules.autoencoding.temporal_ae import VideoDecoder
+from ..modules.diffusionmodules.wrappers import OPENAIUNETWRAPPER
+from ..modules.ema import LitEma
+from ..util import (
+    default,
+    disabled_train,
+    get_obj_from_str,
+    instantiate_from_config,
+    log_txt_as_img,
+)
+
+
+class DiffusionEngine(pl.LightningModule):
+    def __init__(
+        self,
+        network_config,
+        denoiser_config,
+        first_stage_config,
+        conditioner_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        sampler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        optimizer_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        scheduler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        loss_fn_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        network_wrapper: Union[None, str, Dict, ListConfig, OmegaConf] = None,
+        ckpt_path: Union[None, str] = None,
+        remove_keys_from_weights: Union[None, List, Tuple] = None,
+        pattern_to_remove: Union[None, str] = None,
+        remove_keys_from_unet_weights: Union[None, List, Tuple] = None,
+        use_ema: bool = False,
+        ema_decay_rate: float = 0.9999,
+        scale_factor: float = 1.0,
+        disable_first_stage_autocast=False,
+        input_key: str = "jpg",
+        log_keys: Union[List, None] = None,
+        no_log_keys: Union[List, None] = None,
+        no_cond_log: bool = False,
+        compile_model: bool = False,
+        en_and_decode_n_samples_a_time: Optional[int] = None,
+        only_train_ipadapter: Optional[bool] = False,
+        to_unfreeze: Optional[List[str]] = [],
+        to_freeze: Optional[List[str]] = [],
+        separate_unet_ckpt: Optional[str] = None,
+        use_thunder: Optional[bool] = False,
+        is_dubbing: Optional[bool] = False,
+        bad_model_path: Optional[str] = None,
+        bad_model_config: Optional[Dict] = None,
+    ):
+        super().__init__()
+
+        # self.automatic_optimization = False
+        self.log_keys = log_keys
+        self.no_log_keys = no_log_keys
+        self.input_key = input_key
+        self.is_dubbing = is_dubbing
+        self.optimizer_config = default(
+            optimizer_config, {"target": "torch.optim.AdamW"}
+        )
+        self.model = self.initialize_network(
+            network_config, network_wrapper, compile_model=compile_model
+        )
+
+        self.denoiser = instantiate_from_config(denoiser_config)
+
+        self.sampler = (
+            instantiate_from_config(sampler_config)
+            if sampler_config is not None
+            else None
+        )
+        self.is_guided = True
+        if (
+            self.sampler
+            and "IdentityGuider" in sampler_config["params"]["guider_config"]["target"]
+        ):
+            self.is_guided = False
+        if self.sampler is not None:
+            config_guider = sampler_config["params"]["guider_config"]
+            sampler_config["params"]["guider_config"] = None
+            self.sampler_no_guidance = instantiate_from_config(sampler_config)
+            sampler_config["params"]["guider_config"] = config_guider
+        self.conditioner = instantiate_from_config(
+            default(conditioner_config, UNCONDITIONAL_CONFIG)
+        )
+        self.scheduler_config = scheduler_config
+        self._init_first_stage(first_stage_config)
+
+        self.loss_fn = (
+            instantiate_from_config(loss_fn_config)
+            if loss_fn_config is not None
+            else None
+        )
+
+        self.use_ema = use_ema
+        if self.use_ema:
+            self.model_ema = LitEma(self.model, decay=ema_decay_rate)
+            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        self.scale_factor = scale_factor
+        self.disable_first_stage_autocast = disable_first_stage_autocast
+        self.no_cond_log = no_cond_log
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(
+                ckpt_path,
+                remove_keys_from_weights=remove_keys_from_weights,
+                pattern_to_remove=pattern_to_remove,
+            )
+            if separate_unet_ckpt is not None:
+                sd = torch.load(separate_unet_ckpt)["state_dict"]
+                if remove_keys_from_unet_weights is not None:
+                    for k in list(sd.keys()):
+                        for remove_key in remove_keys_from_unet_weights:
+                            if remove_key in k:
+                                del sd[k]
+                self.model.diffusion_model.load_state_dict(sd, strict=False)
+
+        self.en_and_decode_n_samples_a_time = en_and_decode_n_samples_a_time
+        print(
+            "Using",
+            self.en_and_decode_n_samples_a_time,
+            "samples at a time for encoding and decoding",
+        )
+
+        if to_freeze:
+            for name, p in self.model.diffusion_model.named_parameters():
+                for layer in to_freeze:
+                    if layer[0] == "!":
+                        if layer[1:] not in name:
+                            # print("Freezing", name)
+                            p.requires_grad = False
+                    else:
+                        if layer in name:
+                            # print("Freezing", name)
+                            p.requires_grad = False
+                # if "time_" in name:
+                #     print("Freezing", name)
+                #     p.requires_grad = False
+
+        if only_train_ipadapter:
+            # Freeze the model
+            for p in self.model.parameters():
+                p.requires_grad = False
+            # Unfreeze the adapter projection layer
+            for p in self.model.diffusion_model.encoder_hid_proj.parameters():
+                p.requires_grad = True
+            # Unfreeze the cross-attention layer
+            for att_layer in self.model.diffusion_model.attn_processors.values():
+                if isinstance(att_layer, IPAdapterAttnProcessor2_0):
+                    for p in att_layer.parameters():
+                        p.requires_grad = True
+
+            # for name, p in self.named_parameters():
+            #     if p.requires_grad:
+            #         print(name)
+
+        if to_unfreeze:
+            for name in to_unfreeze:
+                for p in getattr(self.model.diffusion_model, name).parameters():
+                    p.requires_grad = True
+
+        if use_thunder:
+            import thunder
+
+            self.model.diffusion_model = thunder.jit(self.model.diffusion_model)
+
+        if "Karras" in denoiser_config.target:
+            assert bad_model_path is not None, (
+                "bad_model_path must be provided for KarrasGuidanceDenoiser"
+            )
+            karras_config = default(bad_model_config, network_config)
+            bad_model = self.initialize_network(
+                karras_config, network_wrapper, compile_model=compile_model
+            )
+            state_dict = self.load_bad_model_weights(bad_model_path)
+            bad_model.load_state_dict(state_dict)
+            self.denoiser.set_bad_network(bad_model)
+
+    def load_bad_model_weights(self, path: str) -> None:
+        print(f"Restoring bad model from {path}")
+        state_dict = torch.load(path, map_location="cpu")
+        new_dict = {}
+        for k, v in state_dict["module"].items():
+            if "learned_mask" in k:
+                new_dict[k.replace("_forward_module.", "").replace("model.", "")] = v
+            if "diffusion_model" in k:
+                new_dict["diffusion_model" + k.split("diffusion_model")[1]] = v
+        return new_dict
+
+    def initialize_network(self, network_config, network_wrapper, compile_model=False):
+        model = instantiate_from_config(network_config)
+        if isinstance(network_wrapper, str) or network_wrapper is None:
+            model = get_obj_from_str(default(network_wrapper, OPENAIUNETWRAPPER))(
+                model, compile_model=compile_model
+            )
+        else:
+            target = network_wrapper["target"]
+            params = network_wrapper.get("params", dict())
+            model = get_obj_from_str(target)(
+                model, compile_model=compile_model, **params
+            )
+        return model
+
+    def init_from_ckpt(
+        self,
+        path: str,
+        remove_keys_from_weights: Optional[Union[List, Tuple]] = None,
+        pattern_to_remove: str = None,
+    ) -> None:
+        print(f"Restoring from {path}")
+        if path.endswith("ckpt"):
+            sd = torch.load(path, map_location="cpu")["state_dict"]
+        elif path.endswith("pt"):
+            sd = torch.load(path, map_location="cpu")["module"]
+            # Remove leading _forward_module from keys
+            sd = {k.replace("_forward_module.", ""): v for k, v in sd.items()}
+        elif path.endswith("bin"):
+            sd = torch.load(path, map_location="cpu")
+            # Remove leading _forward_module from keys
+            sd = {k.replace("_forward_module.", ""): v for k, v in sd.items()}
+        elif path.endswith("safetensors"):
+            sd = load_safetensors(path)
+        else:
+            raise NotImplementedError
+
+        print(f"Loaded state dict from {path} with {len(sd)} keys")
+
+        # if remove_keys_from_weights is not None:
+        #     for k in list(sd.keys()):
+        #         for remove_key in remove_keys_from_weights:
+        #             if remove_key in k:
+        #                 del sd[k]
+        if pattern_to_remove is not None or remove_keys_from_weights is not None:
+            sd = self.remove_mismatched_keys(
+                sd, pattern_to_remove, remove_keys_from_weights
+            )
+
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(
+            f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+        )
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+        if len(unexpected) > 0:
+            print(f"Unexpected Keys: {unexpected}")
+
+    def remove_mismatched_keys(self, state_dict, pattern=None, additional_keys=None):
+        """Remove keys from the state dictionary based on a pattern and a list of additional specific keys."""
+        # Find keys that match the pattern
+        if pattern is not None:
+            mismatched_keys = [key for key in state_dict if re.search(pattern, key)]
+        else:
+            mismatched_keys = []
+
+        print(f"Removing {len(mismatched_keys)} keys based on pattern {pattern}")
+        print(mismatched_keys)
+
+        # Add specific keys to be removed
+        if additional_keys:
+            mismatched_keys.extend(
+                [key for key in additional_keys if key in state_dict]
+            )
+
+        # Remove all identified keys
+        for key in mismatched_keys:
+            if key in state_dict:
+                del state_dict[key]
+        return state_dict
+
+    def _init_first_stage(self, config):
+        model = instantiate_from_config(config).eval()
+        model.train = disabled_train
+        for param in model.parameters():
+            param.requires_grad = False
+        self.first_stage_model = model
+        if self.input_key == "latents":
+            # Remove encoder to save memory
+            self.first_stage_model.encoder = None
+        torch.cuda.empty_cache()
+
+    def get_input(self, batch):
+        # assuming unified data format, dataloader returns a dict.
+        # image tensors should be scaled to -1 ... 1 and in bchw format
+        return batch[self.input_key]
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        is_video = False
+        if len(z.shape) == 5:
+            is_video = True
+            T = z.shape[2]
+            z = rearrange(z, "b c t h w -> (b t) c h w")
+
+        z = 1.0 / self.scale_factor * z
+        n_samples = default(self.en_and_decode_n_samples_a_time, z.shape[0])
+
+        n_rounds = math.ceil(z.shape[0] / n_samples)
+        all_out = []
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            for n in range(n_rounds):
+                if isinstance(self.first_stage_model.decoder, VideoDecoder):
+                    kwargs = {"timesteps": len(z[n * n_samples : (n + 1) * n_samples])}
+                else:
+                    kwargs = {}
+                out = self.first_stage_model.decode(
+                    z[n * n_samples : (n + 1) * n_samples], **kwargs
+                )
+                all_out.append(out)
+        out = torch.cat(all_out, dim=0)
+        if is_video:
+            out = rearrange(out, "(b t) c h w -> b c t h w", t=T)
+        torch.cuda.empty_cache()
+        return out
+
+    @torch.no_grad()
+    def encode_first_stage(self, x):
+        is_video = False
+        if len(x.shape) == 5:
+            is_video = True
+            T = x.shape[2]
+            x = rearrange(x, "b c t h w -> (b t) c h w")
+        n_samples = default(self.en_and_decode_n_samples_a_time, x.shape[0])
+        n_rounds = math.ceil(x.shape[0] / n_samples)
+        all_out = []
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            for n in range(n_rounds):
+                out = self.first_stage_model.encode(
+                    x[n * n_samples : (n + 1) * n_samples]
+                )
+                all_out.append(out)
+        z = torch.cat(all_out, dim=0)
+        z = self.scale_factor * z
+        if is_video:
+            z = rearrange(z, "(b t) c h w -> b c t h w", t=T)
+        return z
+
+    def forward(self, x, batch):
+        loss_dict = self.loss_fn(
+            self.model,
+            self.denoiser,
+            self.conditioner,
+            x,
+            batch,
+            self.first_stage_model,
+        )
+        # loss_mean = loss.mean()
+        for k in loss_dict:
+            loss_dict[k] = loss_dict[k].mean()
+        # loss_dict = {"loss": loss_mean}
+        return loss_dict["loss"], loss_dict
+
+    def shared_step(self, batch: Dict) -> Any:
+        x = self.get_input(batch)
+        if self.input_key != "latents":
+            x = self.encode_first_stage(x)
+        batch["global_step"] = self.global_step
+        loss, loss_dict = self(x, batch)
+        return loss, loss_dict
+
+    def training_step(self, batch, batch_idx):
+        loss, loss_dict = self.shared_step(batch)
+        # debugging_message = "Training step"
+        # print(f"RANK - {self.trainer.global_rank}: {debugging_message}")
+
+        self.log_dict(
+            loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=False
+        )
+
+        self.log(
+            "global_step",
+            self.global_step,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=False,
+        )
+
+        # debugging_message = "Training step - log"
+        # print(f"RANK - {self.trainer.global_rank}: {debugging_message}")
+
+        if self.scheduler_config is not None:
+            lr = self.optimizers().param_groups[0]["lr"]
+            self.log(
+                "lr_abs", lr, prog_bar=True, logger=True, on_step=True, on_epoch=False
+            )
+
+        # # to prevent other processes from moving forward until all processes are in sync
+        # self.trainer.strategy.barrier()
+
+        return loss
+
+    # def validation_step(self, batch, batch_idx):
+    #     # loss, loss_dict = self.shared_step(batch)
+    #     # self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+    #     self.log(
+    #         "global_step",
+    #         self.global_step,
+    #         prog_bar=True,
+    #         logger=True,
+    #         on_step=True,
+    #         on_epoch=False,
+    #     )
+    #     return 0
+
+    # def on_train_epoch_start(self, *args, **kwargs):
+    #     print(f"RANK - {self.trainer.global_rank}: on_train_epoch_start")
+
+    def on_train_start(self, *args, **kwargs):
+        # os.environ["CUDA_VISIBLE_DEVICES"] = str(self.trainer.global_rank)
+        # torch.cuda.set_device(self.trainer.global_rank)
+        # torch.cuda.empty_cache()
+        if self.sampler is None or self.loss_fn is None:
+            raise ValueError("Sampler and loss function need to be set for training.")
+
+    # def on_before_batch_transfer(self, batch, dataloader_idx):
+    #     print(f"RANK - {self.trainer.global_rank}: on_before_batch_transfer - {dataloader_idx}")
+    #     return batch
+
+    # def on_after_batch_transfer(self, batch, dataloader_idx):
+    #     print(f"RANK - {self.trainer.global_rank}: on_after_batch_transfer - {dataloader_idx}")
+    #     return batch
+
+    def on_train_batch_end(self, *args, **kwargs):
+        # print(f"RANK - {self.trainer.global_rank}: on_train_batch_end")
+        if self.use_ema:
+            self.model_ema(self.model)
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.model.parameters())
+            self.model_ema.copy_to(self.model)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.model.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def instantiate_optimizer_from_config(self, params, lr, cfg):
+        return get_obj_from_str(cfg["target"])(
+            params, lr=lr, **cfg.get("params", dict())
+        )
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        params = list(self.model.parameters())
+        for embedder in self.conditioner.embedders:
+            if embedder.is_trainable:
+                params = params + list(embedder.parameters())
+        opt = self.instantiate_optimizer_from_config(params, lr, self.optimizer_config)
+        if self.scheduler_config is not None:
+            scheduler = instantiate_from_config(self.scheduler_config)
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [
+                {
+                    "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule),
+                    "interval": "step",
+                    "frequency": 1,
+                }
+            ]
+            return [opt], scheduler
+        return opt
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: Dict,
+        uc: Union[Dict, None] = None,
+        batch_size: int = 16,
+        shape: Union[None, Tuple, List] = None,
+        **kwargs,
+    ):
+        randn = torch.randn(batch_size, *shape).to(self.device)
+
+        denoiser = lambda input, sigma, c: self.denoiser(
+            self.model, input, sigma, c, **kwargs
+        )
+        samples = self.sampler(denoiser, randn, cond, uc=uc)
+
+        return samples
+
+    @torch.no_grad()
+    def sample_no_guider(
+        self,
+        cond: Dict,
+        uc: Union[Dict, None] = None,
+        batch_size: int = 16,
+        shape: Union[None, Tuple, List] = None,
+        **kwargs,
+    ):
+        randn = torch.randn(batch_size, *shape).to(self.device)
+
+        denoiser = lambda input, sigma, c: self.denoiser(
+            self.model, input, sigma, c, **kwargs
+        )
+        samples = self.sampler_no_guidance(denoiser, randn, cond, uc=uc)
+
+        return samples
+
+    @torch.no_grad()
+    def log_conditionings(self, batch: Dict, n: int) -> Dict:
+        """
+        Defines heuristics to log different conditionings.
+        These can be lists of strings (text-to-image), tensors, ints, ...
+        """
+        image_h, image_w = batch[self.input_key].shape[-2:]
+        log = dict()
+
+        for embedder in self.conditioner.embedders:
+            if (
+                (self.log_keys is None) or (embedder.input_key in self.log_keys)
+            ) and not self.no_cond_log:
+                if embedder.input_key in self.no_log_keys:
+                    continue
+                x = batch[embedder.input_key][:n]
+                if isinstance(x, torch.Tensor):
+                    if x.dim() == 1:
+                        # class-conditional, convert integer to string
+                        x = [str(x[i].item()) for i in range(x.shape[0])]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 4)
+                    elif x.dim() == 2:
+                        # size and crop cond and the like
+                        x = [
+                            "x".join([str(xx) for xx in x[i].tolist()])
+                            for i in range(x.shape[0])
+                        ]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    elif x.dim() == 4:  # already an image
+                        xc = x
+                    elif x.dim() == 5:
+                        xc = torch.cat([x[:, :, i] for i in range(x.shape[2])], dim=-1)
+                    else:
+                        print(x.shape, embedder.input_key)
+                        raise NotImplementedError()
+                elif isinstance(x, (List, ListConfig)):
+                    if isinstance(x[0], str):
+                        # strings
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    else:
+                        raise NotImplementedError()
+                else:
+                    raise NotImplementedError()
+                log[embedder.input_key] = xc
+        return log
+
+    @torch.no_grad()
+    def log_images(
+        self,
+        batch: Dict,
+        N: int = 8,
+        sample: bool = True,
+        ucg_keys: List[str] = None,
+        **kwargs,
+    ) -> Dict:
+        conditioner_input_keys = [e.input_key for e in self.conditioner.embedders]
+        if ucg_keys:
+            assert all(map(lambda x: x in conditioner_input_keys, ucg_keys)), (
+                "Each defined ucg key for sampling must be in the provided conditioner input keys,"
+                f"but we have {ucg_keys} vs. {conditioner_input_keys}"
+            )
+        else:
+            ucg_keys = conditioner_input_keys
+        log = dict()
+
+        x = self.get_input(batch)
+
+        c, uc = self.conditioner.get_unconditional_conditioning(
+            batch,
+            force_uc_zero_embeddings=ucg_keys
+            if len(self.conditioner.embedders) > 0
+            else [],
+        )
+
+        sampling_kwargs = {}
+
+        N = min(x.shape[0], N)
+        x = x.to(self.device)[:N]
+        if self.input_key != "latents":
+            log["inputs"] = x
+            z = self.encode_first_stage(x)
+        else:
+            z = x
+        log["reconstructions"] = self.decode_first_stage(z)
+        log.update(self.log_conditionings(batch, N))
+
+        for k in c:
+            if isinstance(c[k], torch.Tensor):
+                c[k], uc[k] = map(lambda y: y[k][:N].to(self.device), (c, uc))
+
+        if sample:
+            with self.ema_scope("Plotting"):
+                samples = self.sample(
+                    c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs
+                )
+            samples = self.decode_first_stage(samples)
+
+            log["samples"] = samples
+
+            with self.ema_scope("Plotting"):
+                samples = self.sample_no_guider(
+                    c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs
+                )
+            samples = self.decode_first_stage(samples)
+
+            log["samples_no_guidance"] = samples
+        return log
+
+    @torch.no_grad()
+    def log_videos(
+        self,
+        batch: Dict,
+        N: int = 8,
+        sample: bool = True,
+        ucg_keys: List[str] = None,
+        **kwargs,
+    ) -> Dict:
+        # conditioner_input_keys = [e.input_key for e in self.conditioner.embedders]
+        # if ucg_keys:
+        #     assert all(map(lambda x: x in conditioner_input_keys, ucg_keys)), (
+        #         "Each defined ucg key for sampling must be in the provided conditioner input keys,"
+        #         f"but we have {ucg_keys} vs. {conditioner_input_keys}"
+        #     )
+        # else:
+        #     ucg_keys = conditioner_input_keys
+        log = dict()
+        batch_uc = {}
+
+        x = self.get_input(batch)
+        num_frames = x.shape[2]  # assuming bcthw format
+
+        for key in batch.keys():
+            if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+                batch_uc[key] = torch.clone(batch[key])
+
+        c, uc = self.conditioner.get_unconditional_conditioning(
+            batch,
+            batch_uc=batch_uc,
+            force_uc_zero_embeddings=ucg_keys
+            if ucg_keys is not None
+            else [
+                "cond_frames",
+                "cond_frames_without_noise",
+            ],
+        )
+
+        # for k in ["crossattn", "concat"]:
+        #     uc[k] = repeat(uc[k], "b ... -> b t ...", t=num_frames)
+        #     uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=num_frames)
+        #     c[k] = repeat(c[k], "b ... -> b t ...", t=num_frames)
+        #     c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=num_frames)
+
+        sampling_kwargs = {}
+
+        N = min(x.shape[0], N)
+        x = x.to(self.device)[:N]
+
+        if self.input_key != "latents":
+            log["inputs"] = x
+            z = self.encode_first_stage(x)
+        else:
+            z = x
+        log["reconstructions"] = self.decode_first_stage(z)
+        log.update(self.log_conditionings(batch, N))
+
+        if c.get("masks", None) is not None:
+            # Create a mask reconstruction
+            masks = 1 - c["masks"]
+            t = masks.shape[2]
+            masks = rearrange(masks, "b c t h w -> (b t) c h w")
+            target_size = (
+                log["reconstructions"].shape[-2],
+                log["reconstructions"].shape[-1],
+            )
+            masks = torch.nn.functional.interpolate(
+                masks, size=target_size, mode="nearest"
+            )
+            masks = rearrange(masks, "(b t) c h w -> b c t h w", t=t)
+            log["mask_reconstructions"] = log["reconstructions"] * masks
+
+        for k in c:
+            if isinstance(c[k], torch.Tensor):
+                c[k], uc[k] = map(lambda y: y[k][:N].to(self.device), (c, uc))
+            elif isinstance(c[k], list):
+                for i in range(len(c[k])):
+                    c[k][i], uc[k][i] = map(
+                        lambda y: y[k][i][:N].to(self.device), (c, uc)
+                    )
+
+        if sample:
+            n = 2 if self.is_guided else 1
+            # if num_frames == 1:
+            #     sampling_kwargs["image_only_indicator"] = torch.ones(n, num_frames).to(self.device)
+            # else:
+            sampling_kwargs["image_only_indicator"] = torch.zeros(n, num_frames).to(
+                self.device
+            )
+            sampling_kwargs["num_video_frames"] = batch["num_video_frames"]
+
+            with self.ema_scope("Plotting"):
+                samples = self.sample(
+                    c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs
+                )
+            samples = self.decode_first_stage(samples)
+            if self.is_dubbing:
+                samples[:, :, :, : samples.shape[-2] // 2] = log["reconstructions"][
+                    :, :, :, : samples.shape[-2] // 2
+                ]
+            log["samples"] = samples
+
+            # Without guidance
+            # if num_frames == 1:
+            #     sampling_kwargs["image_only_indicator"] = torch.ones(1, num_frames).to(self.device)
+            # else:
+            sampling_kwargs["image_only_indicator"] = torch.zeros(1, num_frames).to(
+                self.device
+            )
+            sampling_kwargs["num_video_frames"] = batch["num_video_frames"]
+
+            with self.ema_scope("Plotting"):
+                samples = self.sample_no_guider(
+                    c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs
+                )
+            samples = self.decode_first_stage(samples)
+            if self.is_dubbing:
+                samples[:, :, :, : samples.shape[-2] // 2] = log["reconstructions"][
+                    :, :, :, : samples.shape[-2] // 2
+                ]
+            log["samples_no_guidance"] = samples
+
+        torch.cuda.empty_cache()
+        return log

sgm/modules/__init__.py

@@ -0,0 +1,6 @@
+from .encoders.modules import GeneralConditioner
+
+UNCONDITIONAL_CONFIG = {
+    "target": "sgm.modules.GeneralConditioner",
+    "params": {"emb_models": []},
+}

sgm/modules/attention.py

@@ -0,0 +1,889 @@
+import logging
+import math
+from inspect import isfunction
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from packaging import version
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+logpy = logging.getLogger(__name__)
+
+if version.parse(torch.__version__) >= version.parse("2.0.0"):
+    SDP_IS_AVAILABLE = True
+    from torch.backends.cuda import SDPBackend, sdp_kernel
+
+    BACKEND_MAP = {
+        SDPBackend.MATH: {
+            "enable_math": True,
+            "enable_flash": False,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.FLASH_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": True,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.EFFICIENT_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": False,
+            "enable_mem_efficient": True,
+        },
+        None: {"enable_math": True, "enable_flash": True, "enable_mem_efficient": True},
+    }
+else:
+    from contextlib import nullcontext
+
+    SDP_IS_AVAILABLE = False
+    sdp_kernel = nullcontext
+    BACKEND_MAP = {}
+    logpy.warn(
+        f"No SDP backend available, likely because you are running in pytorch "
+        f"versions < 2.0. In fact, you are using PyTorch {torch.__version__}. "
+        f"You might want to consider upgrading."
+    )
+
+try:
+    import xformers
+    import xformers.ops
+
+    XFORMERS_IS_AVAILABLE = True
+except:
+    XFORMERS_IS_AVAILABLE = False
+    logpy.warn("no module 'xformers'. Processing without...")
+
+# from .diffusionmodules.util import mixed_checkpoint as checkpoint
+
+
+def exists(val):
+    return val is not None
+
+
+def uniq(arr):
+    return {el: True for el in arr}.keys()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+
+        self.net = nn.Sequential(
+            project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+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 Normalize(in_channels):
+    return torch.nn.GroupNorm(
+        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(
+            qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
+        )
+        k = k.softmax(dim=-1)
+        context = torch.einsum("bhdn,bhen->bhde", k, v)
+        out = torch.einsum("bhde,bhdn->bhen", context, q)
+        out = rearrange(
+            out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
+        )
+        return self.to_out(out)
+
+
+class SelfAttention(nn.Module):
+    ATTENTION_MODES = ("xformers", "torch", "math")
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: Optional[float] = None,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        attn_mode: str = "xformers",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        assert attn_mode in self.ATTENTION_MODES
+        self.attn_mode = attn_mode
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, C = x.shape
+
+        qkv = self.qkv(x)
+        if self.attn_mode == "torch":
+            qkv = rearrange(
+                qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+            ).float()
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B H L D
+            x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+            x = rearrange(x, "B H L D -> B L (H D)")
+        elif self.attn_mode == "xformers":
+            qkv = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B L H D
+            x = xformers.ops.memory_efficient_attention(q, k, v)
+            x = rearrange(x, "B L H D -> B L (H D)", H=self.num_heads)
+        elif self.attn_mode == "math":
+            qkv = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B H L D
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = (attn @ v).transpose(1, 2).reshape(B, L, C)
+        else:
+            raise NotImplemented
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.k = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b (h w) c")
+        k = rearrange(k, "b c h w -> b c (h w)")
+        w_ = torch.einsum("bij,bjk->bik", q, k)
+
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, "b c h w -> b c (h w)")
+        w_ = rearrange(w_, "b i j -> b j i")
+        h_ = torch.einsum("bij,bjk->bik", v, w_)
+        h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim,
+        context_dim=None,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        backend=None,
+        **kwargs,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+        )
+        self.backend = backend
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+        skip_attention=None,
+        **kwargs,
+    ):
+        h = self.heads
+
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+
+        # Ensure skip_attention is a B×1 boolean tensor
+        if skip_attention is None:
+            skip_attention = torch.zeros_like(x[:, :1], dtype=torch.bool)
+
+        assert isinstance(skip_attention, torch.Tensor)
+        assert skip_attention.shape[1] == 1 and skip_attention.dtype == torch.bool
+
+        # Split the batch into skip and non-skip parts
+        skip_indices = skip_attention.squeeze(1)
+        non_skip_indices = ~skip_indices
+
+        # Process skip attention samples
+        if skip_indices.any():
+            x_skip = x[skip_indices]
+            out_skip = self.to_v(x_skip)
+            out_skip = rearrange(out_skip, "b n (h d) -> b n (h d)", h=h)
+
+        # If all samples are skipped, return early
+        if not non_skip_indices.any():
+            if additional_tokens is not None:
+                out_skip = out_skip[:, n_tokens_to_mask:]
+            return self.to_out(out_skip)
+
+        # Process non-skip samples with attention
+        x_non_skip = x[non_skip_indices]
+        q = self.to_q(x_non_skip)
+        context = default(context, x_non_skip)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x_non_skip.shape[0] % n_times_crossframe_attn_in_self == 0
+            k = repeat(
+                k[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+            v = repeat(
+                v[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+        with sdp_kernel(**BACKEND_MAP[self.backend]):
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        out = rearrange(out, "b h n d -> b n (h d)", h=h)
+
+        # Combine skip and non-skip results
+        combined_out = torch.zeros(
+            (x.shape[0], out.shape[1], out.shape[2]), dtype=out.dtype, device=out.device
+        )
+        combined_out[non_skip_indices] = out
+        if skip_indices.any():
+            combined_out[skip_indices] = out_skip
+
+        if additional_tokens is not None:
+            combined_out = combined_out[:, n_tokens_to_mask:]
+        return self.to_out(combined_out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+    def __init__(
+        self,
+        query_dim,
+        context_dim=None,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        use_reference=False,
+        extra_linear=False,
+        **kwargs,
+    ):
+        super().__init__()
+        logpy.debug(
+            f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, "
+            f"context_dim is {context_dim} and using {heads} heads with a "
+            f"dimension of {dim_head}."
+        )
+        inner_dim = dim_head * heads
+        self.is_context = context_dim is not None
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+        self.use_reference = use_reference and self.is_context
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        if not self.use_reference:
+            self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+            self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        else:
+            if extra_linear:
+                self.to_k = nn.Linear(inner_dim, inner_dim, bias=False)
+                self.to_v = nn.Linear(inner_dim, inner_dim, bias=False)
+        self.extra_linear = extra_linear
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+        )
+        self.attention_op: Optional[Any] = None
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+        skip_attention=None,
+    ):
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+
+        # Ensure skip_attention is a B×1 boolean tensor
+        if skip_attention is None:
+            skip_attention = torch.zeros(x.shape[0], 1, dtype=torch.bool)
+        # print(x.shape)
+        # print(skip_attention)
+        # print(skip_attention.shape)
+        # print(any(skip_attention))
+        assert isinstance(skip_attention, torch.Tensor)
+        assert skip_attention.shape[1] == 1 and skip_attention.dtype == torch.bool
+
+        # Split the batch into skip and non-skip parts
+        skip_indices = skip_attention.squeeze(1)
+        non_skip_indices = ~skip_indices
+
+        # Process skip attention samples
+        if skip_indices.any():
+            x_skip = x[skip_indices]
+            out_skip = self.to_v(x_skip)
+            out_skip = (
+                out_skip.unsqueeze(0)
+                .reshape(-1, self.heads, out_skip.shape[1], self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(-1, out_skip.shape[1], self.heads * self.dim_head)
+            )
+        # If all samples are skipped, return early
+        if not non_skip_indices.any():
+            if additional_tokens is not None:
+                out_skip = out_skip[:, n_tokens_to_mask:]
+            return self.to_out(out_skip)
+
+        x_non_skip = x[non_skip_indices]
+        q = self.to_q(x_non_skip)
+        if not self.use_reference:
+            context = default(context, x_non_skip)
+            k = self.to_k(context)
+            v = self.to_v(context)
+        else:
+            # Reference has already correct shape
+            assert context is not None
+            if self.extra_linear:
+                k = self.to_k(context)
+                v = self.to_v(context)
+            k, v = context, context
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x_non_skip.shape[0] % n_times_crossframe_attn_in_self == 0
+            # n_cp = x.shape[0]//n_times_crossframe_attn_in_self
+            k = repeat(
+                k[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+            v = repeat(
+                v[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+        if q.dtype != k.dtype:
+            k = k.to(q.dtype)
+            v = v.to(q.dtype)
+
+        # actually compute the attention, what we cannot get enough of
+        if version.parse(xformers.__version__) >= version.parse("0.0.21"):
+            # NOTE: workaround for
+            # https://github.com/facebookresearch/xformers/issues/845
+            max_bs = 32768
+            N = q.shape[0]
+            n_batches = math.ceil(N / max_bs)
+            out = list()
+            for i_batch in range(n_batches):
+                batch = slice(i_batch * max_bs, (i_batch + 1) * max_bs)
+                out.append(
+                    xformers.ops.memory_efficient_attention(
+                        q[batch],
+                        k[batch],
+                        v[batch],
+                        attn_bias=None,
+                        op=self.attention_op,
+                    )
+                )
+            out = torch.cat(out, 0)
+        else:
+            out = xformers.ops.memory_efficient_attention(
+                q, k, v, attn_bias=None, op=self.attention_op
+            )
+
+        # TODO: Use this directly in the attention operation, as a bias
+        if exists(mask):
+            raise NotImplementedError
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        # Combine skip and non-skip results
+        combined_out = torch.zeros(
+            (x.shape[0], out.shape[1], out.shape[2]), dtype=out.dtype, device=out.device
+        )
+        combined_out[non_skip_indices] = out
+        if skip_indices.any():
+            combined_out[skip_indices] = out_skip
+        else:
+            combined_out = out
+
+        if additional_tokens is not None:
+            # remove additional token
+            combined_out = combined_out[:, n_tokens_to_mask:]
+        return self.to_out(combined_out)
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention,  # ampere
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        disable_self_attn=False,
+        attn_mode="softmax",
+        sdp_backend=None,
+        reference_to=None,
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE:
+            logpy.warn(
+                f"Attention mode '{attn_mode}' is not available. Falling "
+                f"back to native attention. This is not a problem in "
+                f"Pytorch >= 2.0. FYI, you are running with PyTorch "
+                f"version {torch.__version__}."
+            )
+            attn_mode = "softmax"
+        elif attn_mode == "softmax" and not SDP_IS_AVAILABLE:
+            logpy.warn(
+                "We do not support vanilla attention anymore, as it is too "
+                "expensive. Sorry."
+            )
+            if not XFORMERS_IS_AVAILABLE:
+                assert False, (
+                    "Please install xformers via e.g. 'pip install xformers==0.0.16'"
+                )
+            else:
+                logpy.info("Falling back to xformers efficient attention.")
+                attn_mode = "softmax-xformers"
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            assert sdp_backend is None or isinstance(sdp_backend, SDPBackend)
+        else:
+            assert sdp_backend is None
+        self.disable_self_attn = disable_self_attn
+        extra_linear = (reference_to is not None) and ("extra" in reference_to)
+        if extra_linear:
+            reference_to = reference_to.replace("_extra", "")
+        assert reference_to in [None, "self", "cross"]
+        self.reference_to = reference_to
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim
+            if (self.disable_self_attn or reference_to == "self")
+            else None,
+            backend=sdp_backend,
+            use_reference=reference_to == "self",
+            extra_linear=extra_linear,
+        )  # is a self-attention if not self.disable_self_attn
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            backend=sdp_backend,
+            use_reference=reference_to == "cross",
+            extra_linear=extra_linear,
+        )  # is self-attn if context is none
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+        if self.checkpoint:
+            logpy.debug(f"{self.__class__.__name__} is using checkpointing")
+
+    def forward(
+        self,
+        x,
+        context=None,
+        reference_context=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+        skip_attention=None,
+    ):
+        kwargs = {"x": x}
+
+        if context is not None:
+            kwargs.update({"context": context})
+
+        if reference_context is not None:
+            kwargs.update({"reference_context": reference_context})
+
+        if additional_tokens is not None:
+            kwargs.update({"additional_tokens": additional_tokens})
+
+        if n_times_crossframe_attn_in_self:
+            kwargs.update(
+                {"n_times_crossframe_attn_in_self": n_times_crossframe_attn_in_self}
+            )
+
+        # return mixed_checkpoint(self._forward, kwargs, self.parameters(), self.checkpoint)
+        if self.checkpoint:
+            # inputs = {"x": x, "context": context}
+            return checkpoint(
+                self._forward, x, context, reference_context, None, 0, skip_attention
+            )
+            # return checkpoint(self._forward, inputs, self.parameters(), self.checkpoint)
+        else:
+            return self._forward(**kwargs)
+
+    def _forward(
+        self,
+        x,
+        context=None,
+        reference_context=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+        skip_attention=None,
+    ):
+        self_context = reference_context if self.reference_to == "self" else context
+        # print(self.reference_to)
+        # print("context: ", context.shape if context is not None else None)
+        # print("reference_context: ", reference_context.shape if reference_context is not None else None)
+        # print("x: ", x.shape)
+
+        x = (
+            self.attn1(
+                self.norm1(x),
+                context=self_context
+                if (self.disable_self_attn or self.reference_to == "self")
+                else None,
+                additional_tokens=additional_tokens,
+                n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self
+                if not self.disable_self_attn
+                else 0,
+                skip_attention=skip_attention,
+            )
+            + x
+        )
+        cross_context = reference_context if self.reference_to == "cross" else context
+        x = (
+            self.attn2(
+                self.norm2(x),
+                context=cross_context,
+                additional_tokens=additional_tokens,
+            )
+            + x
+        )
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class BasicTransformerSingleLayerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention,  # on the A100s not quite as fast as the above version
+        # (todo might depend on head_dim, check, falls back to semi-optimized kernels for dim!=[16,32,64,128])
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        attn_mode="softmax",
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        # inputs = {"x": x, "context": context}
+        # return checkpoint(self._forward, inputs, self.parameters(), self.checkpoint)
+        return checkpoint(self._forward, x, context)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context) + x
+        x = self.ff(self.norm2(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    NEW: use_linear for more efficiency instead of the 1x1 convs
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        depth=1,
+        dropout=0.0,
+        context_dim=None,
+        disable_self_attn=False,
+        use_linear=False,
+        attn_type="softmax",
+        use_checkpoint=True,
+        # sdp_backend=SDPBackend.FLASH_ATTENTION
+        sdp_backend=None,
+        reference_to=None,
+    ):
+        super().__init__()
+        logpy.debug(
+            f"constructing {self.__class__.__name__} of depth {depth} w/ "
+            f"{in_channels} channels and {n_heads} heads."
+        )
+
+        if exists(context_dim) and not isinstance(context_dim, list):
+            context_dim = [context_dim]
+        if exists(context_dim) and isinstance(context_dim, list):
+            if depth != len(context_dim):
+                logpy.warn(
+                    f"{self.__class__.__name__}: Found context dims "
+                    f"{context_dim} of depth {len(context_dim)}, which does not "
+                    f"match the specified 'depth' of {depth}. Setting context_dim "
+                    f"to {depth * [context_dim[0]]} now."
+                )
+                # depth does not match context dims.
+                assert all(map(lambda x: x == context_dim[0], context_dim)), (
+                    "need homogenous context_dim to match depth automatically"
+                )
+                context_dim = depth * [context_dim[0]]
+        elif context_dim is None:
+            context_dim = [None] * depth
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d],
+                    disable_self_attn=disable_self_attn,
+                    attn_mode=attn_type,
+                    checkpoint=use_checkpoint,
+                    sdp_backend=sdp_backend,
+                    reference_to=reference_to,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(
+                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            )
+        else:
+            # self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None, skip_attention=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            if i > 0 and len(context) == 1:
+                i = 0  # use same context for each block
+            x = block(x, context=context[i], skip_attention=skip_attention)
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in
+
+
+class SimpleTransformer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        heads: int,
+        dim_head: int,
+        context_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        checkpoint: bool = True,
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                BasicTransformerBlock(
+                    dim,
+                    heads,
+                    dim_head,
+                    dropout=dropout,
+                    context_dim=context_dim,
+                    attn_mode="softmax-xformers",
+                    checkpoint=checkpoint,
+                )
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x, context)
+        return x