Upload folder using huggingface_hub

.dockerignore

@@ -0,0 +1,10 @@
+**/__pycache__/
+.deepsource.toml
+.env
+.github/
+.idea/
+.mypy_cache/
+.ruff_cache/
+.venv/
+.vscode/
+renovate.json

.gitignore

@@ -0,0 +1,17 @@
+**/__pycache__/
+**.dccache
+**.env
+.mypy_cache/
+.ruff_cache/
+.venv/
+logs/
+results/
+ckpts/
+Anitalker/
+outputs/
+
+.github/
+.trunk/
+.idea/
+renovate.json
+.deepsource.toml

.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+fail_fast: false
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-yaml
+      - id: check-toml
+      - id: check-json
+      - id: trailing-whitespace
+      - id: check-merge-conflict
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    rev: 0.7.13
+    hooks:
+      - id: uv-lock
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.0
+    hooks:
+      - id: ruff-check
+        args: [ --fix ]
+      - id: ruff-format
+  - repo: https://github.com/abravalheri/validate-pyproject
+    rev: v0.24.1
+    hooks:
+      - id: validate-pyproject

.python-version

@@ -0,0 +1 @@
+3.10

Dockerfile

@@ -0,0 +1,46 @@
+FROM python:3.10 AS builder
+
+SHELL ["/bin/bash", "-c"]
+
+ENV UV_LINK_MODE=copy \
+    UV_COMPILE_BYTECODE=1 \
+    UV_PYTHON_DOWNLOADS=0
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+WORKDIR /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    --mount=type=bind,source=README.md,target=README.md \
+    uv sync --no-install-project --no-dev --locked --no-editable
+
+COPY . /app
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --no-dev --locked --no-editable
+
+FROM python:3.10-slim AS production
+
+SHELL ["/bin/bash", "-c"]
+
+ENV GRADIO_SERVER_PORT=7860 \
+    GRADIO_SERVER_NAME=0.0.0.0
+# skipcq: DOK-DL3008
+RUN groupadd app && \
+    useradd -m -g app -s /bin/bash app && \
+    apt-get update -qq && \
+    apt-get install -qq -y --no-install-recommends espeak-ng ffmpeg && \
+    apt-get clean -qq && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /home/app
+
+COPY --from=builder --chown=app:app /app/.venv /app/.venv
+
+USER app
+
+EXPOSE ${GRADIO_SERVER_PORT}
+
+CMD ["/app/.venv/bin/vocalizr"]

README.md

@@ -1,12 +1,10 @@
 ---
 title: Visualizr
-emoji: 🏆
-colorFrom: green
-colorTo: gray
-sdk: gradio
-sdk_version: 5.34.0
-app_file: app.py
-pinned: false
+emoji: 👁️
+colorFrom: gray
+colorTo: pink
+sdk: docker
+app_port: 7860
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## **Visualizr**: Video Generator part of the Chatacter Backend

compose.yaml

@@ -0,0 +1,26 @@
+name: Visualizr
+services:
+  visualizr:
+    image: ghcr.io/alphaspheredotai/visualizr:latest
+    ports:
+      - "7860:7860"
+    volumes:
+      - visualizr_venv:/venv
+    restart: on-failure:3
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:7860"]
+      interval: 1m30s
+      timeout: 10s
+      retries: 5
+      start_period: 40s
+      start_interval: 5s
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+volumes:
+  visualizr_venv:
+    name: Python Environment for Visualizr

pyproject.toml

@@ -0,0 +1,56 @@
+[project]
+name = "visualizr"
+description = "Video Generator part of the Chatacter Backend"
+version = "0.1.0"
+readme = "README.md"
+requires-python = "~=3.10"
+dependencies = [
+	"espnet>=202412",
+	"gfpgan>=1.3.8",
+	"gradio[mcp]>=5.34.0",
+	"librosa>=0.9.2",
+	"loguru>=0.7.3",
+	"moviepy>=1.0.3",
+	"python-speech-features>=0.6",
+	"pytorch-lightning>=2.5.1.post0",
+	"realesrgan>=0.3.0",
+	"spaces>=0.37.0",
+	"torch>=2.6.0",
+	"torchaudio>=2.6.0",
+	"torchmetrics>=1.7.1",
+	"torchvision>=0.21.0",
+	"tqdm>=4.67.1",
+	"transformers>=4.52.2",
+]
+
+[project.scripts]
+visualizr = "visualizr.__main__:main"
+
+[build-system]
+build-backend = "uv_build"
+requires = ["uv_build"]
+
+[dependency-groups]
+dev = [
+	"huggingface-hub[cli,hf-transfer]>=0.33.0",
+	"pyrefly>=0.19.2",
+	"ruff>=0.11.10",
+	"ty>=0.0.1a6",
+	"watchfiles>=1.0.5",
+]
+
+[tool.uv.sources]
+torch = [
+	{ index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchaudio = [
+	{ index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchvision = [
+	{ index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+[[tool.uv.index]]
+explicit = true
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"

qodana.yaml

@@ -0,0 +1,44 @@
+version: "1.0"
+linter: jetbrains/qodana-python:2025.1
+bootstrap: |
+  export UV_CACHE_DIR=/data/cache
+  export CONDA_PREFIX=/opt/miniconda3
+  mkdir .log
+  conda config --add channels defaults > .log/output.log
+  conda install -y python=3.10 > .log/output.log
+  pip install uv
+  uv pip sync pyproject.toml --python 3.10
+  rm -rd .log
+exclude:
+  - name: All
+    paths:
+      - uv.lock
+profile:
+  name: qodana.recommended
+include:
+  - name: CheckDependencyLicenses
+  - name: PyArgumentListInspection
+  - name: PyTypeCheckerInspection
+  - name: PyDataclassInspection
+  - name: PyUnresolvedReferencesInspection
+  - name: CyclomaticComplexityInspection
+  - name: UnsatisfiedRequirementInspection
+  - name: IgnoreFileDuplicateEntry
+  - name: YAMLSchemaDeprecation
+  - name: YAMLDuplicatedKeys
+  - name: YAMLRecursiveAlias
+  - name: YAMLIncompatibleTypes
+  - name: YAMLUnresolvedAlias
+  - name: YAMLUnusedAnchor
+  - name: YAMLSchemaValidation
+  - name: CheckModuleLicenses
+  - name: CheckThirdPartySoftwareList
+  - name: Annotator
+  - name: EmptyDirectory
+  - name: InconsistentLineSeparators
+  - name: IncorrectFormatting
+  - name: LongLine
+  - name: ProblematicWhitespace
+  - name: ReassignedToPlainText
+  - name: RedundantSuppression
+  - name: TodoComment

src/visualizr/LIA_Model.py

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

src/visualizr/__init__.py

@@ -0,0 +1,59 @@
+from datetime import datetime
+from os import getenv
+from pathlib import Path
+
+from dotenv import load_dotenv
+from huggingface_hub import snapshot_download
+from loguru import logger
+from torch import cuda
+
+load_dotenv()
+
+DEBUG: bool = getenv(key="DEBUG", default="True").lower() == "true"
+SERVER_NAME: str = getenv(key="GRADIO_SERVER_NAME", default="localhost")
+SERVER_PORT: int = int(getenv(key="GRADIO_SERVER_PORT", default="8080"))
+CURRENT_DATE: str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+BASE_DIR: Path = Path.cwd()
+RESULTS_DIR: Path = BASE_DIR / "results"
+LOG_DIR: Path = BASE_DIR / "logs"
+CHECKPOINT_DIR: Path = BASE_DIR / "ckpts"
+AUDIO_FILE_PATH: Path = RESULTS_DIR / f"{CURRENT_DATE}.wav"
+LOG_FILE_PATH: Path = LOG_DIR / f"{CURRENT_DATE}.log"
+CUDA_AVAILABLE: bool = cuda.is_available()
+
+FRAMES_RESULT_SAVED_PATH: Path = RESULTS_DIR / "frames"
+STAGE_1_CHECKPOINT_PATH = CHECKPOINT_DIR / "stage1.ckpt"
+VIDEO_PATH = RESULTS_DIR / f"{CURRENT_DATE}.mp4"
+
+RESULTS_DIR.mkdir(exist_ok=True)
+LOG_DIR.mkdir(exist_ok=True)
+CHECKPOINT_DIR.mkdir(exist_ok=True)
+FRAMES_RESULT_SAVED_PATH.mkdir(exist_ok=True)
+
+MOTION_DIM: int = 20
+TMP_MP4: str = ".tmp.mp4"
+
+logger.add(
+    sink=LOG_FILE_PATH,
+    format="{time:YYYY-MM-DD at HH:mm:ss} | {level} | {message}",
+    colorize=True,
+)
+logger.info(f"CUDA Available: {CUDA_AVAILABLE}")
+logger.info(f"Current date: {CURRENT_DATE}")
+logger.info(f"Base directory: {BASE_DIR}")
+logger.info(f"Results directory: {RESULTS_DIR}")
+logger.info(f"Log directory: {LOG_DIR}")
+logger.info(f"Checkpoint directory: {CHECKPOINT_DIR}")
+
+model_mapping: dict[str, str] = {
+    "mfcc_pose_only": f"{CHECKPOINT_DIR}/stage2_pose_only_mfcc.ckpt",
+    "mfcc_full_control": f"{CHECKPOINT_DIR}/stage2_more_controllable_mfcc.ckpt",
+    "hubert_audio_only": f"{CHECKPOINT_DIR}/stage2_audio_only_hubert.ckpt",
+    "hubert_pose_only": f"{CHECKPOINT_DIR}/stage2_pose_only_hubert.ckpt",
+    "hubert_full_control": f"{CHECKPOINT_DIR}/stage2_full_control_hubert.ckpt",
+}
+
+snapshot_download(
+    repo_id="taocode/anitalker_ckpts", local_dir=CHECKPOINT_DIR, repo_type="model"
+)

src/visualizr/__main__.py

@@ -0,0 +1,22 @@
+from gradio import Blocks
+
+from visualizr import DEBUG, SERVER_NAME, SERVER_PORT
+from visualizr.gui import app_block
+
+
+def main() -> None:
+    """Launch the Gradio voice generation web application."""
+    app: Blocks = app_block()
+    app.queue(api_open=True).launch(
+        server_name=SERVER_NAME,
+        server_port=SERVER_PORT,
+        debug=DEBUG,
+        mcp_server=True,
+        show_api=True,
+        enable_monitoring=True,
+        show_error=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/visualizr/choices.py

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

src/visualizr/config.py

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

src/visualizr/config_base.py

@@ -0,0 +1,83 @@
+import json
+from copy import deepcopy
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class BaseConfig:
+    """BaseConfig provides methods to clone itself,
+    inherit settings from another config, propagate settings to nested configs,
+    and serialize/deserialize configurations to/from JSON.
+    """
+
+    def clone(self):
+        """Return a deep copy of this configuration."""
+        return deepcopy(self)
+
+    def inherit(self, another):
+        """inherit common keys from a given config"""
+        common_keys = set(self.__dict__.keys()) & set(another.__dict__.keys())
+        for k in common_keys:
+            setattr(self, k, getattr(another, k))
+
+    def propagate(self):
+        """push down the configuration to all members"""
+        for _, v in self.__dict__.items():
+            if isinstance(v, BaseConfig):
+                v.inherit(self)
+                v.propagate()
+
+    def save(self, save_path: Path):
+        """save config to JSON file"""
+        if not save_path.exists():
+            save_path.mkdir(parents=True, exist_ok=True)
+        conf = self.as_dict_jsonable()
+        with open(save_path, "w") as f:
+            json.dump(conf, f)
+
+    def load(self, load_path: Path):
+        """load json config"""
+        if not load_path.exists():
+            load_path.mkdir(parents=True, exist_ok=True)
+        with open(load_path) as f:
+            conf = json.load(f)
+        self.from_dict(conf)
+
+    def from_dict(self, config_dict, strict=False):
+        """Populate configuration attributes from a dictionary, optionally
+        enforcing strict key checking.
+        """
+        for k, v in config_dict.items():
+            if not hasattr(self, k):
+                if strict:
+                    raise ValueError(f"loading extra '{k}'")
+                print(f"loading extra '{k}'")
+                continue
+            if isinstance(self.__dict__[k], BaseConfig):
+                self.__dict__[k].from_dict(v)
+            else:
+                self.__dict__[k] = v
+
+    def as_dict_jsonable(self):
+        """Convert the configuration to a JSON-serializable dictionary."""
+        conf = {}
+        for k, v in self.__dict__.items():
+            if isinstance(v, BaseConfig):
+                conf[k] = v.as_dict_jsonable()
+            else:
+                if jsonable(v):
+                    conf[k] = v
+                else:
+                    # ignore not jsonable
+                    pass
+        return conf
+
+
+def jsonable(x):
+    """Check if the object x is JSON serializable."""
+    try:
+        json.dumps(x)
+        return True
+    except TypeError:
+        return False

src/visualizr/dataset.py

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

src/visualizr/dataset_util.py

@@ -0,0 +1,14 @@
+from os import path
+from shutil import copytree
+
+from visualizr import logger
+from visualizr.dist_utils import barrier, get_rank
+
+
+def use_cached_dataset_path(source_path, cache_path):
+    if get_rank() == 0:
+        if not path.exists(cache_path):
+            logger.info(f"copying the data: {source_path} to {cache_path}")
+            copytree(source_path, cache_path)
+    barrier()
+    return cache_path

src/visualizr/diffusion/__init__.py

@@ -0,0 +1,9 @@
+from typing import Union
+
+from visualizr.diffusion.diffusion import (
+    SpacedDiffusionBeatGans,
+    SpacedDiffusionBeatGansConfig,
+)
+
+Sampler = Union[SpacedDiffusionBeatGans]
+SamplerConfig = Union[SpacedDiffusionBeatGansConfig]

src/visualizr/diffusion/base.py

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

src/visualizr/diffusion/diffusion.py

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

src/visualizr/diffusion/resample.py

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

src/visualizr/dist_utils.py

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

src/visualizr/experiment.py

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

src/visualizr/face_sr/face_enhancer.py

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

src/visualizr/face_sr/videoio.py

@@ -0,0 +1,14 @@
+import cv2
+import numpy as np
+
+
+def load_video_to_cv2(input_path: str) -> list[np.ndarray]:
+    video_stream = cv2.VideoCapture(input_path)
+    full_frames = []
+    while 1:
+        still_reading, frame = video_stream.read()
+        if not still_reading:
+            video_stream.release()
+            break
+        full_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    return full_frames

src/visualizr/gui.py

@@ -0,0 +1,97 @@
+from gradio import (
+    Accordion,
+    Audio,
+    Blocks,
+    Button,
+    Checkbox,
+    Column,
+    Dropdown,
+    Image,
+    Markdown,
+    Number,
+    Row,
+    Slider,
+    Video,
+)
+
+from visualizr.settings import DefaultValues
+from visualizr.utils import generate_video
+
+
+def app_block() -> Blocks:
+    """Create the Gradio interface for the voice generation web application."""
+    with Blocks() as app:
+        Markdown(value="# AniTalker")
+        with Row():
+            with Column():
+                uploaded_img: Image = Image(type="filepath", label="Reference Image")
+                uploaded_audio = Audio(
+                    type="filepath", label="Input Audio", show_download_button=True
+                )
+            with Column():
+                output_video_256 = Video(label="Generated Video (256)")
+                output_video_512 = Video(label="Generated Video (512)")
+                output_message = Markdown()
+
+        generate_button = Button(value="Generate Video")
+
+        with Accordion(label="Configuration"):
+            infer_type = Dropdown(
+                label="Inference Type",
+                choices=[
+                    "mfcc_full_control",
+                    "mfcc_pose_only",
+                    "hubert_pose_only",
+                    "hubert_audio_only",
+                    "hubert_full_control",
+                ],
+                value="hubert_audio_only",
+            )
+            face_sr = Checkbox(label="Enable Face Super-Resolution (512*512)")
+            seed = Number(label="Seed", value=DefaultValues().seed)
+            pose_yaw = Slider(
+                label="pose_yaw",
+                minimum=-1,
+                maximum=1,
+                value=DefaultValues().pose_yaw,
+            )
+            pose_pitch = Slider(
+                label="pose_pitch",
+                minimum=-1,
+                maximum=1,
+                value=DefaultValues().pose_pitch,
+            )
+            pose_roll = Slider(
+                label="pose_roll",
+                minimum=-1,
+                maximum=1,
+                value=DefaultValues().pose_roll,
+            )
+            face_location = Slider(
+                label="face_location", maximum=1, value=DefaultValues().face_location
+            )
+            face_scale = Slider(
+                label="face_scale", maximum=1, value=DefaultValues().face_scale
+            )
+            step_t = Slider(
+                label="step_T", minimum=1, step=1, value=DefaultValues().step_T
+            )
+
+        generate_button.click(
+            fn=generate_video,
+            inputs=[
+                uploaded_img,
+                uploaded_audio,
+                infer_type,
+                pose_yaw,
+                pose_pitch,
+                pose_roll,
+                face_location,
+                face_scale,
+                step_t,
+                face_sr,
+                seed,
+            ],
+            outputs=[output_video_256, output_video_512, output_message],
+        )
+        return app

src/visualizr/model/__init__.py

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

src/visualizr/model/base.py

@@ -0,0 +1,28 @@
+import numpy as np
+import torch
+
+
+class BaseModule(torch.nn.Module):
+    def __init__(self):
+        super(BaseModule, self).__init__()
+
+    @property
+    def nparams(self):
+        """
+        Returns number of trainable parameters of the module.
+        """
+        num_params = 0
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                num_params += np.prod(param.detach().cpu().numpy().shape)
+        return num_params
+
+    def relocate_input(self, x: list):
+        """
+        Relocates provided tensors to the same device set for the module.
+        """
+        device = next(self.parameters()).device
+        for i in range(len(x)):
+            if isinstance(x[i], torch.Tensor) and x[i].device != device:
+                x[i] = x[i].to(device)
+        return x

src/visualizr/model/blocks.py

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

src/visualizr/model/diffusion.py

@@ -0,0 +1,323 @@
+import math
+
+import torch
+from einops import rearrange
+
+from visualizr.model.base import BaseModule
+
+
+class Mish(BaseModule):
+    def forward(self, x):
+        return x * torch.tanh(torch.nn.functional.softplus(x))
+
+
+class Upsample(BaseModule):
+    def __init__(self, dim):
+        super(Upsample, self).__init__()
+        self.conv = torch.nn.ConvTranspose2d(dim, dim, 4, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Downsample(BaseModule):
+    def __init__(self, dim):
+        super(Downsample, self).__init__()
+        self.conv = torch.nn.Conv2d(dim, dim, 3, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Rezero(BaseModule):
+    def __init__(self, fn):
+        super(Rezero, self).__init__()
+        self.fn = fn
+        self.g = torch.nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        return self.fn(x) * self.g
+
+
+class Block(BaseModule):
+    def __init__(self, dim, dim_out, groups=8):
+        super(Block, self).__init__()
+        self.block = torch.nn.Sequential(
+            torch.nn.Conv2d(dim, dim_out, 3, padding=1),
+            torch.nn.GroupNorm(groups, dim_out),
+            Mish(),
+        )
+
+    def forward(self, x, mask):
+        output = self.block(x * mask)
+        return output * mask
+
+
+class ResnetBlock(BaseModule):
+    def __init__(self, dim, dim_out, time_emb_dim, groups=8):
+        super(ResnetBlock, self).__init__()
+        self.mlp = torch.nn.Sequential(Mish(), torch.nn.Linear(time_emb_dim, dim_out))
+
+        self.block1 = Block(dim, dim_out, groups=groups)
+        self.block2 = Block(dim_out, dim_out, groups=groups)
+        if dim != dim_out:
+            self.res_conv = torch.nn.Conv2d(dim, dim_out, 1)
+        else:
+            self.res_conv = torch.nn.Identity()
+
+    def forward(self, x, mask, time_emb):
+        h = self.block1(x, mask)
+        h += self.mlp(time_emb).unsqueeze(-1).unsqueeze(-1)
+        h = self.block2(h, mask)
+        output = h + self.res_conv(x * mask)
+        return output
+
+
+class LinearAttention(BaseModule):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super(LinearAttention, self).__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = torch.nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = torch.nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(
+            qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
+        )
+        k = k.softmax(dim=-1)
+        context = torch.einsum("bhdn,bhen->bhde", k, v)
+        out = torch.einsum("bhde,bhdn->bhen", context, q)
+        out = rearrange(
+            out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
+        )
+        return self.to_out(out)
+
+
+class Residual(BaseModule):
+    def __init__(self, fn):
+        super(Residual, self).__init__()
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        output = self.fn(x, *args, **kwargs) + x
+        return output
+
+
+class SinusoidalPosEmb(BaseModule):
+    def __init__(self, dim):
+        super(SinusoidalPosEmb, self).__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class GradLogPEstimator2d(BaseModule):
+    def __init__(
+        self,
+        dim,
+        dim_mults=(1, 2, 4),
+        groups=8,
+        n_spks=None,
+        spk_emb_dim=64,
+        n_feats=80,
+        pe_scale=1000,
+    ):
+        super(GradLogPEstimator2d, self).__init__()
+        self.dim = dim
+        self.dim_mults = dim_mults
+        self.groups = groups
+        self.n_spks = n_spks if not isinstance(n_spks, type(None)) else 1
+        self.spk_emb_dim = spk_emb_dim
+        self.pe_scale = pe_scale
+
+        if n_spks > 1:
+            self.spk_mlp = torch.nn.Sequential(
+                torch.nn.Linear(spk_emb_dim, spk_emb_dim * 4),
+                Mish(),
+                torch.nn.Linear(spk_emb_dim * 4, n_feats),
+            )
+        self.time_pos_emb = SinusoidalPosEmb(dim)
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(dim, dim * 4), Mish(), torch.nn.Linear(dim * 4, dim)
+        )
+
+        dims = [2 + (1 if n_spks > 1 else 0), *map(lambda m: dim * m, dim_mults)]
+        in_out = list(zip(dims[:-1], dims[1:]))
+        self.downs = torch.nn.ModuleList([])
+        self.ups = torch.nn.ModuleList([])
+        num_resolutions = len(in_out)
+
+        for ind, (dim_in, dim_out) in enumerate(in_out):
+            is_last = ind >= (num_resolutions - 1)
+            self.downs.append(
+                torch.nn.ModuleList(
+                    [
+                        ResnetBlock(dim_in, dim_out, time_emb_dim=dim),
+                        ResnetBlock(dim_out, dim_out, time_emb_dim=dim),
+                        Residual(Rezero(LinearAttention(dim_out))),
+                        Downsample(dim_out) if not is_last else torch.nn.Identity(),
+                    ]
+                )
+            )
+
+        mid_dim = dims[-1]
+        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)
+        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
+        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)
+
+        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
+            self.ups.append(
+                torch.nn.ModuleList(
+                    [
+                        ResnetBlock(dim_out * 2, dim_in, time_emb_dim=dim),
+                        ResnetBlock(dim_in, dim_in, time_emb_dim=dim),
+                        Residual(Rezero(LinearAttention(dim_in))),
+                        Upsample(dim_in),
+                    ]
+                )
+            )
+        self.final_block = Block(dim, dim)
+        self.final_conv = torch.nn.Conv2d(dim, 1, 1)
+
+    def forward(self, x, mask, mu, t, spk=None):
+        global s
+        if not isinstance(spk, type(None)):
+            s = self.spk_mlp(spk)
+
+        t = self.time_pos_emb(t, scale=self.pe_scale)
+        t = self.mlp(t)
+
+        if self.n_spks < 2:
+            x = torch.stack([mu, x], 1)
+        else:
+            s = s.unsqueeze(-1).repeat(1, 1, x.shape[-1])
+            x = torch.stack([mu, x, s], 1)
+        mask = mask.unsqueeze(1)
+
+        hiddens = []
+        masks = [mask]
+        for resnet1, resnet2, attn, downsample in self.downs:
+            mask_down = masks[-1]
+            x = resnet1(x, mask_down, t)
+            x = resnet2(x, mask_down, t)
+            x = attn(x)
+            hiddens.append(x)
+            x = downsample(x * mask_down)
+            masks.append(mask_down[:, :, :, ::2])
+
+        masks = masks[:-1]
+        mask_mid = masks[-1]
+        x = self.mid_block1(x, mask_mid, t)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, mask_mid, t)
+
+        for resnet1, resnet2, attn, upsample in self.ups:
+            mask_up = masks.pop()
+            x = torch.cat((x, hiddens.pop()), dim=1)
+            x = resnet1(x, mask_up, t)
+            x = resnet2(x, mask_up, t)
+            x = attn(x)
+            x = upsample(x * mask_up)
+
+        x = self.final_block(x, mask)
+        output = self.final_conv(x * mask)
+
+        return (output * mask).squeeze(1)
+
+
+def get_noise(t, beta_init, beta_term, cumulative=False):
+    if cumulative:
+        noise = beta_init * t + 0.5 * (beta_term - beta_init) * (t**2)
+    else:
+        noise = beta_init + (beta_term - beta_init) * t
+    return noise
+
+
+class Diffusion(BaseModule):
+    def __init__(
+        self,
+        n_feats,
+        dim,
+        n_spks=1,
+        spk_emb_dim=64,
+        beta_min=0.05,
+        beta_max=20,
+        pe_scale=1000,
+    ):
+        super(Diffusion, self).__init__()
+        self.n_feats = n_feats
+        self.dim = dim
+        self.n_spks = n_spks
+        self.spk_emb_dim = spk_emb_dim
+        self.beta_min = beta_min
+        self.beta_max = beta_max
+        self.pe_scale = pe_scale
+
+        self.estimator = GradLogPEstimator2d(
+            dim, n_spks=n_spks, spk_emb_dim=spk_emb_dim, pe_scale=pe_scale
+        )
+
+    def forward_diffusion(self, x0, mask, mu, t):
+        time = t.unsqueeze(-1).unsqueeze(-1)
+        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True)
+        mean = x0 * torch.exp(-0.5 * cum_noise) + mu * (
+            1.0 - torch.exp(-0.5 * cum_noise)
+        )
+        variance = 1.0 - torch.exp(-cum_noise)
+        z = torch.randn(x0.shape, dtype=x0.dtype, device=x0.device, requires_grad=False)
+        xt = mean + z * torch.sqrt(variance)
+        return xt * mask, z * mask
+
+    @torch.no_grad()
+    def reverse_diffusion(self, z, mask, mu, n_timesteps, stoc=False, spk=None):
+        h = 1.0 / n_timesteps
+        xt = z * mask
+        for i in range(n_timesteps):
+            t = (1.0 - (i + 0.5) * h) * torch.ones(
+                z.shape[0], dtype=z.dtype, device=z.device
+            )
+            time = t.unsqueeze(-1).unsqueeze(-1)
+            noise_t = get_noise(time, self.beta_min, self.beta_max, cumulative=False)
+            if stoc:  # adds stochastic term
+                dxt_det = 0.5 * (mu - xt) - self.estimator(xt, mask, mu, t, spk)
+                dxt_det = dxt_det * noise_t * h
+                dxt_stoc = torch.randn(
+                    z.shape, dtype=z.dtype, device=z.device, requires_grad=False
+                )
+                dxt_stoc = dxt_stoc * torch.sqrt(noise_t * h)
+                dxt = dxt_det + dxt_stoc
+            else:
+                dxt = 0.5 * (mu - xt - self.estimator(xt, mask, mu, t, spk))
+                dxt = dxt * noise_t * h
+            xt = (xt - dxt) * mask
+        return xt
+
+    @torch.no_grad()
+    def forward(self, z, mask, mu, n_timesteps, stoc=False, spk=None):
+        return self.reverse_diffusion(z, mask, mu, n_timesteps, stoc, spk)
+
+    def loss_t(self, x0, mask, mu, t, spk=None):
+        xt, z = self.forward_diffusion(x0, mask, mu, t)
+        time = t.unsqueeze(-1).unsqueeze(-1)
+        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True)
+        noise_estimation = self.estimator(xt, mask, mu, t, spk)
+        noise_estimation *= torch.sqrt(1.0 - torch.exp(-cum_noise))
+        loss = torch.sum((noise_estimation + z) ** 2) / (torch.sum(mask) * self.n_feats)
+        return loss, xt
+
+    def compute_loss(self, x0, mask, mu, spk=None, offset=1e-5):
+        t = torch.rand(
+            x0.shape[0], dtype=x0.dtype, device=x0.device, requires_grad=False
+        )
+        t = torch.clamp(t, offset, 1.0 - offset)
+        return self.loss_t(x0, mask, mu, t, spk)

src/visualizr/model/latentnet.py

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

src/visualizr/model/nn.py

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

src/visualizr/model/seq2seq.py

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

src/visualizr/model/unet.py

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

src/visualizr/model/unet_autoenc.py

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

src/visualizr/networks/discriminator.py

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

src/visualizr/networks/encoder.py

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

src/visualizr/networks/generator.py

@@ -0,0 +1,37 @@
+from torch import nn
+
+from visualizr.networks.encoder import Encoder
+from visualizr.networks.styledecoder import Synthesis
+
+
+class Generator(nn.Module):
+    def __init__(
+        self,
+        size,
+        style_dim=512,
+        motion_dim=20,
+        channel_multiplier=1,
+        blur_kernel=[1, 3, 3, 1],
+    ):
+        super(Generator, self).__init__()
+
+        # encoder
+        self.enc = Encoder(size, style_dim, motion_dim)
+        self.dec = Synthesis(
+            size, style_dim, motion_dim, blur_kernel, channel_multiplier
+        )
+
+    def get_direction(self):
+        return self.dec.direction(None)
+
+    def synthesis(self, wa, alpha, feat):
+        img = self.dec(wa, alpha, feat)
+
+        return img
+
+    def forward(self, img_source, img_drive, h_start=None):
+        wa, alpha, feats = self.enc(img_source, img_drive, h_start)
+        # import pdb;pdb.set_trace()
+        img_recon = self.dec(wa, alpha, feats)
+
+        return img_recon

src/visualizr/networks/styledecoder.py

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

src/visualizr/networks/utils.py

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

src/visualizr/renderer.py

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

src/visualizr/settings.py

@@ -0,0 +1,29 @@
+from typing import Literal
+
+from pydantic import BaseModel
+from torch.cuda import is_available
+
+
+class Args(BaseModel):
+    test_image_path: str
+    test_audio_path: str
+    test_hubert_path: str
+    result_path: str = "./outputs/"
+    stage1_checkpoint_path: str = "ckpts/stage1.ckpt"
+    stage2_checkpoint_path: str
+    control_flag: bool = True
+    pose_driven_path: str = "not_supported_in_this_mode"
+    image_size: int = 256
+    device: Literal["cuda", "cpu"] = "cuda" if is_available() else "cpu"
+    motion_dim: int = 20
+    decoder_layers: int = 2
+
+
+class DefaultValues(BaseModel):
+    pose_yaw: float = 0.0
+    pose_pitch: float = 0.0
+    pose_roll: float = 0.0
+    face_location: float = 0.5
+    face_scale: float = 0.5
+    step_T: int = 50
+    seed: int = 0

src/visualizr/templates.py

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

src/visualizr/utils.py

@@ -0,0 +1,432 @@
+import os
+import shutil
+import sys
+import time
+from importlib.util import find_spec
+from pathlib import Path
+from typing import Literal
+
+import gradio as gr
+import librosa
+import numpy as np
+import python_speech_features
+import spaces
+import torch
+from gradio import Markdown
+from moviepy.editor import (
+    AudioFileClip,
+    ImageClip,
+    VideoFileClip,
+    concatenate_videoclips,
+)
+from PIL import Image
+from torch import Tensor
+from torchvision.transforms import ToPILImage
+from tqdm import tqdm
+
+from visualizr import (
+    FRAMES_RESULT_SAVED_PATH,
+    MOTION_DIM,
+    RESULTS_DIR,
+    STAGE_1_CHECKPOINT_PATH,
+    TMP_MP4,
+    logger,
+    model_mapping,
+)
+from visualizr.config import TrainConfig
+from visualizr.experiment import LitModel
+from visualizr.LIA_Model import LIA_Model
+from visualizr.templates import ffhq256_autoenc
+
+
+def check_package_installed(package_name: str) -> bool:
+    return find_spec(package_name) is not None
+
+
+def frames_to_video(input_path, audio_path, output_path, fps=25):
+    image_files = [
+        os.path.join(input_path, img) for img in sorted(os.listdir(input_path))
+    ]
+    clips = [ImageClip(m).set_duration(1 / fps) for m in image_files]
+    video = concatenate_videoclips(clips, method="compose")
+    audio = AudioFileClip(audio_path)
+    final_video = video.set_audio(audio)
+    final_video.write_videofile(output_path, fps, "libx264", audio_codec="aac")
+
+
+def load_image(filename: str, size: int) -> np.ndarray:
+    img: Image.Image = Image.open(filename).convert("RGB")
+    img_resized: Image.Image = img.resize((size, size))
+    img_np: np.ndarray = np.asarray(img_resized)
+    img_transposed: np.ndarray = np.transpose(img_np, (2, 0, 1))  # 3 x 256 x 256
+    return img_transposed / 255.0
+
+
+def img_preprocessing(img_path: str, size: int) -> Tensor:
+    img_np: np.ndarray = load_image(img_path, size)  # [0, 1]
+    img: Tensor = torch.from_numpy(img_np).unsqueeze(0).float()  # [0, 1]
+    normalized_image: Tensor = (img - 0.5) * 2.0  # [-1, 1]
+    return normalized_image
+
+
+def saved_image(img_tensor: Tensor, img_path: str) -> None:
+    pil_image_converter: ToPILImage = ToPILImage()
+    img = pil_image_converter(img_tensor.detach().cpu().squeeze(0))
+    img.save(img_path)
+
+
+def load_stage_1_model() -> LIA_Model:
+    logger.info("Loading stage 1 model... ")
+    lia: LIA_Model = LIA_Model(motion_dim=MOTION_DIM, fusion_type="weighted_sum")
+    lia.load_lightning_model(STAGE_1_CHECKPOINT_PATH)
+    lia.to("cuda")
+    return lia
+
+
+def load_stage_2_model(conf: TrainConfig, stage2_checkpoint_path: str) -> LitModel:
+    logger.info("Loading stage 2 model... ")
+    model = LitModel(conf)
+    state = torch.load(stage2_checkpoint_path, "cpu")
+    model.load_state_dict(state)
+    model.ema_model.eval()
+    model.ema_model.to("cuda")
+    return model
+
+
+def init_conf(
+    infer_type: Literal[
+        "mfcc_full_control",
+        "mfcc_pose_only",
+        "hubert_pose_only",
+        "hubert_audio_only",
+        "hubert_full_control",
+    ],
+    seed: int,
+) -> TrainConfig:
+    logger.info("Initializing configuration... ")
+    conf: TrainConfig = ffhq256_autoenc()
+    conf.seed = seed
+    conf.decoder_layers = 2
+    conf.infer_type = infer_type
+    conf.motion_dim = MOTION_DIM
+    logger.info(f"infer_type: {infer_type}")
+    match infer_type:
+        case "mfcc_full_control":
+            conf.face_location = True
+            conf.face_scale = True
+            conf.mfcc = True
+        case "mfcc_pose_only":
+            conf.face_location = False
+            conf.face_scale = False
+            conf.mfcc = True
+        case "hubert_pose_only":
+            conf.face_location = False
+            conf.face_scale = False
+            conf.mfcc = False
+        case "hubert_audio_only":
+            conf.face_location = False
+            conf.face_scale = False
+            conf.mfcc = False
+        case "hubert_full_control":
+            conf.face_location = True
+            conf.face_scale = True
+            conf.mfcc = False
+    return conf
+
+
+def main(
+    infer_type: Literal[
+        "mfcc_full_control",
+        "mfcc_pose_only",
+        "hubert_pose_only",
+        "hubert_audio_only",
+        "hubert_full_control",
+    ],
+    image_path: str,
+    test_audio_path: str,
+    face_sr: bool,
+    pose_yaw: float,
+    pose_pitch: float,
+    pose_roll: float,
+    face_location: float,
+    face_scale: float,
+    step_t: int,
+    seed: int,
+    stage2_checkpoint_path: str,
+):
+    global frame_end, audio_driven
+    if not os.path.exists(image_path):
+        logger.exception(f"{image_path} does not exist!")
+        sys.exit(0)
+    if not os.path.exists(test_audio_path):
+        logger.exception(f"{test_audio_path} does not exist!")
+        sys.exit(0)
+
+    image_name: str = Path(image_path).stem
+    audio_name: str = Path(test_audio_path).stem
+
+    predicted_video_256_path: Path = RESULTS_DIR / f"{image_name}-{audio_name}.mp4"
+    predicted_video_512_path: Path = RESULTS_DIR / f"{image_name}-{audio_name}_SR.mp4"
+
+    # ======Loading Stage 1 model=========
+    lia: LIA_Model = load_stage_1_model()
+    # ============================
+
+    conf: TrainConfig = init_conf(infer_type, seed)
+
+    img_source: Tensor = img_preprocessing(image_path, 256).to("cuda")
+    one_shot_lia_start, one_shot_lia_direction, feats = lia.get_start_direction_code(
+        img_source, img_source, img_source, img_source
+    )
+
+    # ======Loading Stage 2 model=========
+    model = load_stage_2_model(conf, stage2_checkpoint_path)
+    # =================================
+
+    # ======Audio Input=========
+    if conf.infer_type.startswith("mfcc"):
+        # MFCC features
+        wav, sr = librosa.load(test_audio_path, sr=16000)
+        input_values = python_speech_features.mfcc(
+            signal=wav, samplerate=sr, numcep=13, winlen=0.025, winstep=0.01
+        )
+        d_mfcc_feat = python_speech_features.base.delta(input_values, 1)
+        d_mfcc_feat2 = python_speech_features.base.delta(input_values, 2)
+        audio_driven_obj: np.ndarray = np.hstack(
+            (input_values, d_mfcc_feat, d_mfcc_feat2)
+        )
+        frame_start, frame_end = 0, int(audio_driven_obj.shape[0] / 4)
+        audio_start, audio_end = (
+            int(frame_start * 4),
+            int(frame_end * 4),
+        )  # The video frame is fixed to 25 hz, and the audio is fixed to 100 hz
+
+        audio_driven = (
+            torch.Tensor(audio_driven_obj[audio_start:audio_end, :])
+            .unsqueeze(0)
+            .float()
+            .to("cuda")
+        )
+
+    elif conf.infer_type.startswith("hubert"):
+        # Hubert features
+        if not check_package_installed("transformers"):
+            logger.exception("Please install transformers module first.")
+            sys.exit(0)
+        hubert_model_path = "ckpts/chinese-hubert-large"
+        if not os.path.exists(hubert_model_path):
+            logger.exception(
+                "Please download the hubert weight into the ckpts path first."
+            )
+            sys.exit(0)
+        logger.info(
+            "You did not extract the audio features in advance, "
+            + "extracting online now, which will increase processing delay"
+        )
+
+        start_time = time.time()
+
+        # load hubert model
+        from transformers import HubertModel, Wav2Vec2FeatureExtractor
+
+        audio_model = HubertModel.from_pretrained(hubert_model_path).to("cuda")
+        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_path)
+        audio_model.feature_extractor._freeze_parameters()  # skipcq: PYL-W0212
+        audio_model.eval()
+
+        # hubert model forward pass
+        audio, sr = librosa.load(test_audio_path, sr=16000)
+        input_values = feature_extractor(
+            audio,
+            sampling_rate=16000,
+            padding=True,
+            do_normalize=True,
+            return_tensors="pt",
+        ).input_values
+        input_values = input_values.to("cuda")
+        ws_feats = []
+        with torch.no_grad():
+            outputs = audio_model(input_values, output_hidden_states=True)
+            for i in range(len(outputs.hidden_states)):
+                ws_feats.append(outputs.hidden_states[i].detach().cpu().numpy())
+            ws_feat_obj = np.array(ws_feats)
+            ws_feat_obj = np.squeeze(ws_feat_obj, 1)
+            ws_feat_obj = np.pad(
+                ws_feat_obj, ((0, 0), (0, 1), (0, 0)), "edge"
+            )  # align the audio length with the video frame
+
+        execution_time = time.time() - start_time
+        logger.info(f"Extraction Audio Feature: {execution_time:.2f} Seconds")
+
+        audio_driven_obj = ws_feat_obj
+
+        frame_start, frame_end = 0, int(audio_driven_obj.shape[1] / 2)
+        audio_start, audio_end = (
+            int(frame_start * 2),
+            int(frame_end * 2),
+        )  # The video frame is fixed to 25 hz, and the audio is fixed to 50 hz
+
+        audio_driven = (
+            torch.Tensor(audio_driven_obj[:, audio_start:audio_end, :])
+            .unsqueeze(0)
+            .float()
+            .to("cuda")
+        )
+    # ============================
+
+    # Diffusion Noise
+    noisy_t = torch.randn((1, frame_end, MOTION_DIM)).to("cuda")
+
+    # ======Inputs for Attribute Control=========
+    yaw_signal = torch.zeros(1, frame_end, 1).to("cuda") + pose_yaw
+    pitch_signal = torch.zeros(1, frame_end, 1).to("cuda") + pose_pitch
+    roll_signal = torch.zeros(1, frame_end, 1).to("cuda") + pose_roll
+    pose_signal = torch.cat((yaw_signal, pitch_signal, roll_signal), dim=-1)
+
+    pose_signal = torch.clamp(pose_signal, -1, 1)
+
+    face_location_signal = torch.zeros(1, frame_end, 1).to("cuda") + face_location
+    face_scale_tensor = torch.zeros(1, frame_end, 1).to("cuda") + face_scale
+    # ===========================================
+    start_time = time.time()
+    # ======Diffusion De-nosing Process=========
+    generated_directions = model.render(
+        one_shot_lia_start,
+        one_shot_lia_direction,
+        audio_driven,
+        face_location_signal,
+        face_scale_tensor,
+        pose_signal,
+        noisy_t,
+        step_t,
+        True,
+    )
+    # =========================================
+
+    execution_time = time.time() - start_time
+    logger.info(f"Motion Diffusion Model: {execution_time:.2f} Seconds")
+
+    generated_directions = generated_directions.detach().cpu().numpy()
+
+    start_time = time.time()
+    # ======Rendering images frame-by-frame=========
+    for pred_index in tqdm(range(generated_directions.shape[1])):
+        ori_img_recon = lia.render(
+            one_shot_lia_start,
+            torch.Tensor(generated_directions[:, pred_index, :]).to("cuda"),
+            feats,
+        )
+        ori_img_recon = ori_img_recon.clamp(-1, 1)
+        wav_pred = (ori_img_recon.detach() + 1) / 2
+        saved_image(
+            wav_pred, os.path.join(FRAMES_RESULT_SAVED_PATH, f"{pred_index:06d}.png")
+        )
+    # ==============================================
+
+    execution_time = time.time() - start_time
+    logger.info(f"Renderer Model: {execution_time:.2f} Seconds")
+    logger.info(f"Saving video at {predicted_video_256_path}")
+
+    frames_to_video(
+        str(FRAMES_RESULT_SAVED_PATH),
+        test_audio_path,
+        str(predicted_video_256_path),
+    )
+
+    shutil.rmtree(FRAMES_RESULT_SAVED_PATH)
+
+    # Enhancer
+    if face_sr and check_package_installed("gfpgan"):
+        from imageio import mimsave
+
+        from visualizr.face_sr.face_enhancer import enhancer_list
+
+        # Super-resolution
+        mimsave(
+            predicted_video_512_path / TMP_MP4,
+            enhancer_list(predicted_video_256_path, bg_upsampler=None),
+            fps=25.0,
+        )
+
+        # Merge audio and video
+        video_clip = VideoFileClip(predicted_video_512_path / TMP_MP4)
+        audio_clip = AudioFileClip(predicted_video_256_path)
+        final_clip = video_clip.set_audio(audio_clip)
+        final_clip.write_videofile(
+            predicted_video_512_path, codec="libx264", audio_codec="aac"
+        )
+
+        os.remove(predicted_video_512_path / TMP_MP4)
+
+    if face_sr:
+        return predicted_video_256_path, predicted_video_512_path
+    return predicted_video_256_path, predicted_video_256_path
+
+
+@spaces.GPU(duration=300)
+def generate_video(
+    uploaded_img: str,
+    uploaded_audio: str,
+    infer_type: Literal[
+        "mfcc_full_control",
+        "mfcc_pose_only",
+        "hubert_pose_only",
+        "hubert_audio_only",
+        "hubert_full_control",
+    ],
+    pose_yaw: float,
+    pose_pitch: float,
+    pose_roll: float,
+    face_location: float,
+    face_scale: float,
+    step_t: int,
+    face_sr: bool,
+    seed: int,
+):
+    if not uploaded_img or not uploaded_audio:
+        return None, Markdown(
+            "Error: Input image or audio file is empty. "
+            + "Please check and upload both files."
+        )
+    try:
+        output_256_video_path, output_512_video_path = main(
+            infer_type,
+            uploaded_img,
+            uploaded_audio,
+            face_sr,
+            pose_yaw,
+            pose_pitch,
+            pose_roll,
+            face_location,
+            face_scale,
+            step_t,
+            seed,
+            model_mapping.get(
+                infer_type,
+                "default_checkpoint.ckpt",
+            ),
+        )
+
+        if not os.path.exists(output_256_video_path):
+            return None, gr.Markdown(
+                "Error: Video generation failed. "
+                + "Please check your inputs and try again."
+            )
+        if output_256_video_path == output_512_video_path:
+            return (
+                gr.Video(value=output_256_video_path),
+                None,
+                gr.Markdown("Video (256*256 only) generated successfully!"),
+            )
+        return (
+            gr.Video(value=output_256_video_path),
+            gr.Video(value=output_512_video_path),
+            gr.Markdown("Video generated successfully!"),
+        )
+
+    except Exception as e:
+        return (
+            None,
+            None,
+            gr.Markdown(f"Error: An unexpected error occurred - {str(e)}"),
+        )