AnimeFace dataset generation by BFNs
import torch
import torch.nn as nn
from tqdm.auto import tqdm

class BFNContinuousData(nn.Module):
    def __init__(self, unet, in_channels=3, sigma=0.001):
      super(BFNContinuousData, self).__init__()
      self.sigma = torch.tensor(sigma)
      self.in_channels = in_channels
      self.unet = unet
    def forward(self, theta, t, ema=None):
        """
        Forward pass of the Bayesian Flow Network.

        Parameters
        ----------
        theta : torch.Tensor
            Tensor of shape (B, D).
        t : torch.Tensor
            Tensor of shape (B,).

        Returns
        -------
        torch.Tensor
            Output tensor of shape (B, D).
        """
        t = self.get_time_embedding(t)
        if ema is None:
            output = self.unet(theta, t)
        else:
            with ema.average_parameters():
                output = self.unet(theta, t)
        return output

    def get_time_embedding(self, timestep):
        # timestep: (bs, )
        # output: (bs, 320)
        freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32, device=timestep.device) / 160)
        x = timestep[:, None] * freqs[None]
        return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)

    def cts_output_prediction(self, mu, t, gamma, ema=None, t_min=1e-10, x_range=(-1 , 1)):
        x_pred = self.forward(mu, t, ema=ema)
        mask = t < t_min
        gamma = torch.where(mask[:, None, None, None], torch.ones_like(gamma), gamma)
        x_pred = mu / gamma - ((1 - gamma) / gamma).sqrt() * x_pred
        x_pred = torch.clamp(x_pred, x_range[0], x_range[1])
        # print(x_pred)
        return x_pred

    def process_infinity(self, x, t=None, return_params=False):
        """
        x : torch.Tensor
            Tensor of shape (B, C, H, W).
        t : torch.Tensor
            Tensor of shape (B, ).
        Returns
        -------
        torch.Tensor
            Output tensor of shape (1, ).
        """
        bs, c, h, w = x.shape
        if t is None:
            t = torch.rand((bs, ), device=x.device)
        gamma = 1 - self.sigma ** (2 * t[:, None, None, None])  # (B, 1, 1, 1)
        mu = gamma * x + (gamma * (1 - gamma)).sqrt() * torch.randn_like(x) # (B, C, H, W)
        pred = self.cts_output_prediction(mu, t, gamma)
        loss_infinity = -torch.log(self.sigma) * self.sigma ** (-2 * t[:, None, None, None]) * (x - pred)**2
        if return_params:
            return loss_infinity.mean(), mu, pred, t
        else:
            return loss_infinity.mean()
        
    def process_discrete(self, x, step=None, max_step=1000, return_params=False):
        """
        x : torch.Tensor
            Tensor of shape (B, C, H, W).
        t : torch.Tensor
            Tensor of shape (B, ).
        n : int
            number of steps.
        Returns
        -------
        torch.Tensor
            Output tensor of shape (1, ).
        """
        bs, c, h, w = x.shape
        if step is None:
            step = torch.randint(1, max_step + 1, (bs, ), device=x.device)
        t = (step - 1) / max_step
        # t=0　⇨ gamma=0, mu=0, pred=0
        gamma = 1 - self.sigma ** (2 * t[:, None, None, None])
        mu = gamma * x + (gamma * (1 - gamma)).sqrt() * torch.randn_like(x)
        pred = self.cts_output_prediction(mu, t, gamma)
        loss_discrete = max_step * (1 - self.sigma ** (2 / max_step)) / (2 * self.sigma ** (2 * step[:, None, None, None] / max_step)) * (x - pred)**2
        if return_params:
            return loss_discrete.mean(), mu, pred, t
        else:
            return loss_discrete.mean()

    @torch.inference_mode()
    def sample(self, h, w, batch_size=4, steps=1000, ema=None, return_samples=None, device="cpu"):
        self.eval()
        mu = torch.zeros((batch_size, self.in_channels, h, w), device=device)
        rho = 1
        ret_list = []
        for istep in tqdm(range(1, steps+1)):
            t = torch.full((batch_size, ), (istep-1)/steps, device=device)
            x_pred = self.cts_output_prediction(mu, t, 1 - self.sigma ** (2 * t[:, None, None, None]), ema=ema)
            alpha = self.sigma ** (-2*istep/steps) * (1 - self.sigma ** (2/steps))
            y = x_pred + torch.randn_like(x_pred, device=device) / alpha.sqrt()
            mu = (rho * mu + alpha * y) / (rho + alpha)
            rho += alpha
            if return_samples is not None and (istep-1) % return_samples == 0:
                ret_list.append(x_pred)
                
        t = torch.ones((batch_size, ), device=device)
        x_pred = self.cts_output_prediction(mu, t, 1 - self.sigma ** (2 * t[:, None, None, None]), ema=ema)
        if return_samples is None:
            return x_pred
        ret_list.append(x_pred)
        return x_pred, ret_list
bfns/bfn_continuous.py
yoimers/bayesian-flow-networks-animeface
import torch
import torch.nn as nn
from tqdm.auto import tqdm
import math

class Discretised():
    # K: number of Bins
    # k: {1, 2, ..., K} 1-index
    def __init__(self, K: int):
        self.K = K
    def center(self, k: int):
        return (2 * k - 1) / self.K - 1
    def left(self, k: int):
        return self.center(k) - 1 / self.K
    def right(self, k: int):
        return self.center(k) + 1 / self.K
    def find_nearest_index(self, x):
        # x: (B, C, H, W)
        values = (2 * torch.arange(1, self.K+1, device=x.device) - 1) / self.K - 1 # values: (K)
        differences = torch.abs(values.view(1, 1, 1, 1, self.K) - x.unsqueeze(-1))
        return torch.argmin(differences, dim=-1) + 1
    def get_center(self, x):
        nearest_index = self.find_nearest_index(x)
        return self.center(nearest_index)

class BFNDiscretisedData(nn.Module):
    def __init__(self, unet, K: int, in_channels=3, sigma=math.sqrt(0.001), device='cuda'):
        super(BFNDiscretisedData, self).__init__()
        self.sigma = torch.tensor(sigma, device=device)
        self.in_channels = in_channels
        self.unet = unet.to(device)
        self.K = K
        self.discretised = Discretised(K)
        self.Kl = self.discretised.left(torch.arange(1, self.K+1).to(device)) # (K, )
        self.Kr = self.discretised.right(torch.arange(1, self.K+1).to(device))
        self.Kc = self.discretised.center(torch.arange(1, self.K+1).to(device))
    def forward(self, theta, t, ema=None):
        """
        Forward pass of the Bayesian Flow Network.

        Parameters
        ----------
        theta : torch.Tensor
            Tensor of shape (B, C, H, W).
        t : torch.Tensor
            Tensor of shape (B,).

        Returns
        -------
        torch.Tensor
            Output tensor of shape (B, 2*C, H, W).
        """
        t = self.get_time_embedding(t)
        if ema is None:
            output = self.unet(theta, t)
        else:
            with ema.average_parameters():
                output = self.unet(theta, t)
        return output

    def get_time_embedding(self, timestep):
        # timestep: (bs, )
        # output: (bs, 320)
        freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32, device=timestep.device) / 160)
        x = timestep[:, None] * freqs[None]
        return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)

    def discretised_cdf(self, x, mu, sigma):
        """
        Parameters ( arbitrary shape X )
        ----------
        x: torch.Tensor
          Tensor of shape X
        mu: torch.Tensor
          Tensor of shape X
        sigma: torch.Tensor
          Tensor of shape X
        Return
        ------
        torch.Tensor
            Output tensor of shape X
        """
        g = 0.5 * (1 + torch.erf((x - mu) / (sigma * math.sqrt(2))))
        g = torch.where(x >= 1, torch.ones_like(x), g)
        g = torch.where(x <= -1, torch.zeros_like(x), g)
        return g
    
    def discretised_output_distribution(self, mu, t, gamma, ema=None, t_min=1e-10):
        """
        Parameters
        ----------
        mu: torch.Tensor
          Tensor of shape (B, C, H, W)
        t: torch.Tensor
          Tensor of shape (B, )
        gamma: torch.Tensor
          Tensor of shape (B, )
        Return
        ------
        torch.Tensor
            Output tensor of shape (B, C, H, W, K)
        """
        output = self.forward(mu, t, ema=ema)
        _, c, _, _ = output.shape
        mu_eps, ln_sigma_eps = output[:, :c//2, :, :], output[:, c//2:, :, :]
        mask = t < t_min                                                              # want t=0 ⇨ mux=0 sigmax=1
        gamma = torch.where(mask[:, None, None, None], torch.ones_like(gamma), gamma) # t=0 ⇨ gammma=1 ⇨ gam=0 mu=0 mux=0 sigmax=1 
        gam = ((1 - gamma) / gamma).sqrt()
        mux = mu / gamma - gam * mu_eps
        sigmax = torch.where(mask[:, None, None, None], torch.ones_like(ln_sigma_eps), gam * torch.exp(ln_sigma_eps))
        output = self.discretised_cdf(self.Kr.reshape(1, 1, 1, 1, -1), mux.unsqueeze(-1), sigmax.unsqueeze(-1)) - self.discretised_cdf(self.Kl.reshape(1, 1, 1, 1, -1), mux.unsqueeze(-1), sigmax.unsqueeze(-1))
        return output
    
    def process_infinity(self, x, t=None, return_params=False):
        """
        x : torch.Tensor
            Tensor of shape (B, C, H, W).
        t : torch.Tensor
            Tensor of shape (B, ).
        Returns
        -------
        torch.Tensor
            Output tensor of shape (1, ).
        """
        x = self.discretised.get_center(x)
        bs, c, h, w = x.shape
        if t is None:
            t = torch.rand((bs, ), device=x.device)
        gamma = 1 - self.sigma ** (2 * t[:, None, None, None])
        mu = gamma * x + (gamma * (1 - gamma)).sqrt() * torch.randn_like(x)
        p_output = self.discretised_output_distribution(mu, t, gamma) #(B, C, H, W, K)
        k_hat = torch.matmul(p_output, self.Kc) #(B, C, H, W)
        L_infinity = -self.sigma ** (-2 * t[:, None, None, None]) * torch.log(self.sigma) * (x - k_hat)**2
        if return_params:
            return L_infinity.mean(), mu, self.discretised.get_center(k_hat), t
        else:
            return L_infinity.mean()
        
    def process_discrete(self, x, step=None, max_step=1000, return_params=False):
        """
        x : torch.Tensor
            Tensor of shape (B, C, H, W).
        t : torch.Tensor
            Tensor of shape (B, ).
        Returns
        -------
        torch.Tensor
            Output tensor of shape (1, ).
        """
        x = self.discretised.get_center(x)
        bs, c, h, w = x.shape
        if step is None:
            step = torch.randint(1, max_step+1, (bs, ), device=x.device)
        t = (step - 1) / max_step
        gamma = 1 - self.sigma ** (2 * t[:, None, None, None])
        mu = gamma * x + (gamma * (1 - gamma)).sqrt() * torch.randn_like(x)
        p_output = self.discretised_output_distribution(mu, t, gamma) #(B, C, H, W, K)
        k_hat = torch.matmul(p_output, self.Kc) #(B, C, H, W)
        alpha = self.sigma ** (-2 * step[:, None, None, None] / max_step) * (1 - self.sigma ** (2 / max_step))
        y = x + torch.randn_like(x) / alpha.sqrt()
        L_discrete  = -alpha / 2 * (y - x) ** 2 
        L_discrete -= (p_output * (-alpha.unsqueeze(-1) / 2 * (y.unsqueeze(-1) - self.Kc.reshape(1, 1, 1, 1, -1)) ** 2).exp()).sum(-1).log()
        L_discrete *= max_step
        if return_params:
            return L_discrete.mean(), mu, self.discretised.get_center(k_hat), t
        else:
            return L_discrete.mean()
        

    @torch.inference_mode()
    def sample(self, h, w, batch_size=4, steps=1000, ema=None, return_samples=None, device="cpu"):
        self.eval()
        mu = torch.zeros((batch_size, self.in_channels, h, w), device=device)
        rho = 1
        ret_list = []
        for step in tqdm(range(1, steps+1)):
            t = (step - 1) / steps * torch.ones((batch_size, ), device=device)
            k = self.discretised_output_distribution(mu, t, 1 - self.sigma ** (2 * t[:, None, None, None]), ema=ema)
            k = torch.matmul(k, self.Kc)
            k = self.discretised.get_center(k)
            alpha = self.sigma ** (-2 * step / steps) * (1 - self.sigma ** (2 / steps))
            y = k + torch.randn_like(k) / alpha.sqrt()
            mu = (rho * mu + alpha * y) / (rho + alpha)
            rho += alpha
            if return_samples is not None and (step-1) % return_samples == 0:
                ret_list.append(k)
        t = torch.ones((batch_size, ), device=device)
        k = self.discretised_output_distribution(mu, t, 1 - self.sigma ** (2 * t[:, None, None, None]), ema=ema)
        k = torch.matmul(k, self.Kc)
        k = self.discretised.get_center(k)
        if return_samples is None:
            return k
        ret_list.append(k)
        return k, ret_list
bfns/bfn_discretised.py
yoimers/bayesian-flow-networks-animeface


import argparse
from matplotlib import pyplot as plt
import numpy as np

import torch
import torch_ema
from bfns.bfn_continuous import BFNContinuousData
from bfns.bfn_discretised import BFNDiscretisedData
from torchvision.transforms import ToPILImage
from matplotlib.animation import FuncAnimation

from networks.unet import UNet
from train_common import BFNType, TimeType
from utils import default_transform, set_seed
import torchvision.datasets as datasets

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def input_output_image(args, bfn: BFNContinuousData, bfnType: BFNType, timeType: TimeType):
    # Comparison of Input and Output images    
    transform = default_transform(args.height, args.width)
    dataset = datasets.ImageFolder(root="./animeface_valid", transform=transform)
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch, shuffle=True)
    x, _ = next(iter(data_loader))

    mus = [[] for _ in range(args.batch)]
    preds = [[] for _ in range(args.batch)]
    ts = []
    col = 20
    with torch.no_grad():
        for t in np.linspace(0.001, 0.66, col):
            if timeType == TimeType.ContinuousTimeLoss:
                time = torch.full((len(x), ), t, device=device)
                _, mu, pred, _ = bfn.process_infinity(x=x.to(device), t=time, return_params=True)
            elif timeType == TimeType.DiscreteTimeLoss:
                step = torch.full((len(x), ), int(args.step*t), device=device)
                _, mu, pred, _ = bfn.process_discrete(x=x.to(device), step=step, max_step=args.step, return_params=True)
            else:
                raise ValueError("The timeType must be either ContinuousTimeLoss or DiscreteTimeLoss.")
            
            for b in range(len(x)):
                mus[b].append(mu.permute(0, 2, 3, 1).cpu().detach().numpy()[b])
                preds[b].append(pred.permute(0, 2, 3, 1).cpu().detach().numpy()[b])
            ts.append(t)

    fig, ax = plt.subplots(2 * args.batch, col, figsize=(12, args.batch))

    for b in range(args.batch):
        for i in range(col):
            ax[2*b, i].imshow(mus[b][i])
            ax[2*b+1, i].imshow(preds[b][i])
            ax[2*b, i].axis('off')
            ax[2*b+1, i].axis('off')
    for i in range(col):
        ax[0, i].set_title(np.round(ts[i], 3))
    plt.subplots_adjust(left=0.02, right=0.98, top=0.9, bottom=0.02)
    plt.savefig('./outputs/validation_image_plots.png')

def common_generate(args: argparse.ArgumentParser, bfnType: BFNType, timeType: TimeType):
    set_seed(args.seed)
    if bfnType == BFNType.Continuous:
        unet = UNet(3, 3).to(device=device)
    elif bfnType == BFNType.Discretized:
        unet = UNet(3, 6).to(device=device)
    else:
        raise ValueError("The BFNType must be either Continuous or Discretized.")
    unet.load_state_dict(torch.load(args.load_model_path))
    unet.eval()
    
    if bfnType == BFNType.Continuous:
        bfn = BFNContinuousData(unet, in_channels=3, sigma=args.sigma).to(device)
    elif bfnType == BFNType.Discretized:
        bfn = BFNDiscretisedData(unet, K=args.K, in_channels=3, sigma=args.sigma).to(device)
    else:
        raise ValueError("The BFNType must be either Continuous or Discretized.")
    bfn.eval()
    
    ema = torch_ema.ExponentialMovingAverage(unet.parameters(), decay=0.9999)
    ema.to(device)
    
    if args.gif_generate:
        return_samples = args.gif_save_every_n_step
    else:
        return_samples = args.step//10
        
    with torch.no_grad():
        x_hat, out_list = bfn.sample(args.height, args.width, device=device, steps=args.step, return_samples=return_samples, batch_size=args.batch, ema=ema)
    
    if args.gif_generate:
        fig, ax = plt.subplots(1, args.batch, figsize=(2 * args.batch, 2))
        def animate(frame):
            for i in range(args.batch):
                ax[i].clear()
                ax[i].imshow(((out_list[frame][i]+1)/2*255).clamp(0, 255).to(torch.uint8).permute(1, 2, 0).cpu().detach().numpy())
                ax[i].axis('off')
                ax[i].set_title(f'Step {max(1, args.gif_save_every_n_step * frame)}', fontsize=12)
                
        anim = FuncAnimation(fig, animate, frames=args.step//args.gif_save_every_n_step, interval=100)
        plt.subplots_adjust(left=0.02, right=0.98, top=0.9, bottom=0.02)
        anim.save("./outputs/generate_output.gif", writer='pillow')
    else:
        fig, ax = plt.subplots(args.batch, len(out_list), figsize=(8, args.batch))
        for i in range(len(out_list)):
            for b in range(len(out_list[i])):
                ax[b, i].imshow(((out_list[i][b]+1)/2*255).clamp(0, 255).to(torch.uint8).permute(1, 2, 0).cpu().detach().numpy())
                ax[b, i].axis('off')
        for i in range(len(out_list)):
            ax[0, i].set_title(np.round(return_samples*i, 2))
        plt.subplots_adjust(left=0.02, right=0.98, top=0.95, bottom=0.02)
        plt.savefig('./outputs/generate_plots.png') # Generated image every step//10
        
    # Save PIL Image
    to_pil = ToPILImage()
    pil_images = [to_pil(((image+1)/2*255).clamp(0, 255).to(torch.uint8)) for image in x_hat]
    for i, pil_image in enumerate(pil_images):
        pil_image.convert("RGB").save(f'./outputs/image_{i}.png') # Generated images
        
    # Comparison of Input and Output images
    if args.generate_input_output_image:
        input_output_image(args, bfn, bfnType=bfnType, timeType=timeType)

def setup_generate_common_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser.add_argument("--load_model_path", type=str, required=True, default="./models/model.pth", help="Path to the pre-trained model file.")
    parser.add_argument("--data_path", type=str, default="./animeface", help="Path to the dataset directory.")
    parser.add_argument("--batch", type=int, default=4, help="Batch size for generation.")
    parser.add_argument("--step", type=int, default=1000, help="Number of generation steps.")
    parser.add_argument("--sigma", type=float, default=0.001, help="Value of sigma for generation.")
    parser.add_argument("--height", type=int, default=32, help="Height of the generated images.")
    parser.add_argument("--width", type=int, default=32, help="Width of the generated images.")
    parser.add_argument("--K", type=int, default=16, help="Number of Bins.")
    parser.add_argument("--seed", type=int)
    parser.add_argument("--generate_input_output_image", action="store_true", help="Whether to generate input(θ, mu, ...) and output images.")
    parser.add_argument("-gif", "--gif_generate", action="store_true", help="Whether to generate output images gif.")
    parser.add_argument("--gif_save_every_n_step", type=int, default=50)
    return parser
generate_common.py
yoimers/bayesian-flow-networks-animeface
import torch
from generate_common import common_generate, setup_generate_common_parser
from train_common import BFNType, TimeType
import argparse

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

bfnType = BFNType.Continuous
timeType = TimeType.DiscreteTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_generate_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    common_generate(args, bfnType = bfnType, timeType=timeType)
generate_continuours.py
yoimers/bayesian-flow-networks-animeface
import argparse
import torch
from generate_common import common_generate, setup_generate_common_parser
from train_common import BFNType, TimeType

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

bfnType = BFNType.Discretized
timeType = TimeType.ContinuousTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_generate_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    common_generate(args, bfnType = bfnType, timeType=timeType)
generate_discretized.py
yoimers/bayesian-flow-networks-animeface
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import math

# Reference
# https://github.com/kjsman/stable-diffusion-pytorch

class SelfAttention(nn.Module):
    def __init__(self, n_heads, d_embed, in_proj_bias=True, out_proj_bias=True):
        super().__init__()
        self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)
        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
        self.n_heads = n_heads
        self.d_head = d_embed // n_heads

    def forward(self, x, causal_mask=False):
        input_shape = x.shape
        batch_size, sequence_length, d_embed = input_shape
        interim_shape = (batch_size, sequence_length, self.n_heads, self.d_head)

        q, k, v = self.in_proj(x).chunk(3, dim=-1)

        q = q.view(interim_shape).transpose(1, 2)
        k = k.view(interim_shape).transpose(1, 2)
        v = v.view(interim_shape).transpose(1, 2)

        weight = q @ k.transpose(-1, -2)
        if causal_mask:
            mask = torch.ones_like(weight, dtype=torch.bool).triu(1)
            weight.masked_fill_(mask, -torch.inf)
        weight /= math.sqrt(self.d_head)
        weight = F.softmax(weight, dim=-1)

        output = weight @ v
        output = output.transpose(1, 2)
        output = output.reshape(input_shape)
        output = self.out_proj(output)
        return output

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, n_time=1280):
        super().__init__()
        self.groupnorm_feature = nn.GroupNorm(32, in_channels)
        self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.linear_time = nn.Linear(n_time, out_channels)

        self.groupnorm_merged = nn.GroupNorm(32, out_channels)
        self.conv_merged = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)

        self.residual_layer = nn.Identity()
        if in_channels != out_channels:
            self.residual_layer = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)

    def forward(self, feature, time):
        residue = feature

        feature = self.groupnorm_feature(feature)
        feature = F.silu(feature)
        feature = self.conv_feature(feature)

        time = F.silu(time)
        time = self.linear_time(time)

        merged = feature + time.unsqueeze(-1).unsqueeze(-1)
        merged = self.groupnorm_merged(merged)
        merged = F.silu(merged)
        merged = self.conv_merged(merged)
        return merged + self.residual_layer(residue)

class AttentionBlock(nn.Module):
    def __init__(self, n_head: int, n_embd: int):
        super().__init__()
        channels = n_head * n_embd

        self.groupnorm = nn.GroupNorm(32, channels, eps=1e-6)
        self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)

        self.layernorm_1 = nn.LayerNorm(channels)
        self.attention_1 = SelfAttention(n_head, channels, in_proj_bias=False)
        self.layernorm_2 = nn.LayerNorm(channels)

        self.layernorm_3 = nn.LayerNorm(channels)
        self.linear_geglu_1 = nn.Linear(channels, 8 * channels)
        self.linear_geglu_2 = nn.Linear(4 * channels, channels)

        self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)

    def forward(self, x):
        residue_long = x
        x = self.groupnorm(x)
        x = self.conv_input(x)

        n, c, h, w = x.shape
        x = x.view((n, c, h * w)).transpose(-1, -2)

        residue_short = x
        x = self.layernorm_1(x)
        x = self.attention_1(x)
        x += residue_short

        xresidue_short = x
        x = self.layernorm_3(x)
        x, gate = self.linear_geglu_1(x).chunk(2, dim=-1)
        x = x * F.gelu(gate)
        x = self.linear_geglu_2(x)
        x += residue_short

        x = x.transpose(-1, -2).view((n, c, h, w))
        return self.conv_output(x) + residue_long

class Upsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
    def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        return self.conv(x)

class TimeEmbedding(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
        self.linear_2 = nn.Linear(4 * n_embd, 4 * n_embd)

    def forward(self, x):
        x = self.linear_1(x)
        x = F.silu(x)
        x = self.linear_2(x)
        return x

class FinalLayer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.groupnorm = nn.GroupNorm(32, in_channels)
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
    
    def forward(self, x):
        x = self.groupnorm(x)
        x = F.silu(x)
        x = self.conv(x)
        return x
        
class Switch(nn.Sequential):
    def forward(self, x, time):
        for layer in self:
            if isinstance(layer, ResidualBlock):
                x = layer(x, time)
            else:
                x = layer(x)
        return x
    
    
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        # input (3, 64, 64)
        self.encoders = nn.ModuleList([
            Switch(nn.Conv2d(in_channels, 160, kernel_size=3, padding=1)),                # input (160, 64, 64)
            Switch(ResidualBlock(160, 160), AttentionBlock(8, 20)),             # input (160, 64, 64)
            Switch(ResidualBlock(160, 160), AttentionBlock(8, 20)),             # input (160, 64, 64)
            Switch(nn.Conv2d(160, 160, kernel_size=3, stride=2, padding=1)),    # input (160, 32, 32)
            Switch(ResidualBlock(160, 320), AttentionBlock(8, 40)),             # input (320, 32, 32)
            Switch(ResidualBlock(320, 320), AttentionBlock(8, 40)),             # input (320, 32, 32)
            Switch(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),    # input (320, 32, 32)
            # Switch(ResidualBlock(640, 1280), AttentionBlock(8, 160)),           # input (1280, 32, 32)
            # Switch(ResidualBlock(1280, 1280), AttentionBlock(8, 160)),          # input (1280, 32, 32)
            # Switch(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),  # input (1280, 16, 16)
            # Switch(ResidualBlock(1280, 1280)),
            # Switch(ResidualBlock(1280, 1280)),
        ])

        self.bottleneck = Switch(ResidualBlock(320, 320), AttentionBlock(8, 40), ResidualBlock(320, 320)) # input (1280, 16, 16)

        self.decoders = nn.ModuleList([
            # Switch(ResidualBlock(2560, 1280)),
            # Switch(ResidualBlock(2560, 1280)),
            # Switch(ResidualBlock(2560, 1280), Upsample(1280)),                  # input (1280, 16, 16)
            # Switch(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),          # input (1280, 16, 16)
            # Switch(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),          # input (1280, 16, 16)
            Switch(ResidualBlock(640, 640), AttentionBlock(8, 80), Upsample(640)),# input (1280, 16, 16)
            Switch(ResidualBlock(960, 640), AttentionBlock(8, 80)),            # input (1280, 16, 16)
            Switch(ResidualBlock(960, 640), AttentionBlock(8, 80)),
            Switch(ResidualBlock(800, 640), AttentionBlock(8, 80), Upsample(640)),
            Switch(ResidualBlock(800, 320), AttentionBlock(8, 40)),
            Switch(ResidualBlock(480, 320), AttentionBlock(8, 40)),
            Switch(ResidualBlock(480, 320), AttentionBlock(8, 40)),
        ])
        self.time_embedding = TimeEmbedding(320)
        self.final_layer = FinalLayer(320, out_channels)

    def forward(self, x, time):
        time = self.time_embedding(time)
        skip_connections = []
        for layers in self.encoders:
            x = layers(x, time)
            skip_connections.append(x)

        x = self.bottleneck(x, time)
        for layers in self.decoders:
            x = torch.cat((x, skip_connections.pop()), dim=1)
            x = layers(x, time)
        x = self.final_layer(x)
        return x
    
    


if __name__ == '__main__':
    from torchinfo import summary
    # (3, 6) 110,884,646
    # (3, 3) 110,884,646
    test = UNet(3, 3) 
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    test.to(device)
    print(summary(test, input_size=[(1, 3, 64, 64), (1, 320)]))
networks/unet.py
yoimers/bayesian-flow-networks-animeface
from datasets import Dataset, load_dataset
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torchinfo import summary
import math
from torch.optim import AdamW
from torch.optim.lr_scheduler import ExponentialLR
from tqdm.auto import tqdm
import torch_ema
import argparse

import matplotlib.pyplot as plt

import os

import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from bfns.bfn_continuous import BFNContinuousData
from bfns.bfn_discretised import BFNDiscretisedData
from networks.unet import UNet
from utils import default_transform
import argparse
from enum import Enum, auto
from torch.utils.tensorboard import SummaryWriter

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class BFNType(Enum):
    Continuous = auto()
    Discretized = auto()

class TimeType(Enum):
    ContinuousTimeLoss = auto()
    DiscreteTimeLoss = auto()

def train(args: argparse.ArgumentParser, bfnType: BFNType, timeType: TimeType):
    transform = default_transform(args.height, args.width)
    dataset = datasets.ImageFolder(root=args.data_path, transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch, shuffle=True, num_workers=8)
    
    if bfnType == BFNType.Continuous:
        unet = UNet(3, 3).to(device)
    elif bfnType == BFNType.Discretized:
        unet = UNet(3, 6).to(device)
    else:
        raise ValueError("The BFNType must be either Continuous or Discretized.")
    
    if args.load_model_path is not None:
        unet.load_state_dict(torch.load(args.load_model_path))
        
    if bfnType == BFNType.Continuous:
        bfn = BFNContinuousData(unet, sigma=args.sigma).to(device)
    elif bfnType == BFNType.Discretized:
        bfn = BFNDiscretisedData(unet, K=args.K, sigma=args.sigma).to(device)
    else:
        raise ValueError("The BFNType must be either Continuous or Discretized.")
    
    optimizer = AdamW(unet.parameters(), lr=args.lr)
    scheduler = ExponentialLR(optimizer, gamma=0.9)
    ema = torch_ema.ExponentialMovingAverage(unet.parameters(), decay=0.9999)
    ema.to(device)
    writer = SummaryWriter()
    
    num = 1
    for epoch in tqdm(range(1, args.epoch+1), desc='Training', unit='epoch'):
        losses = []
        for X, _ in tqdm(train_loader, desc='Epoch {}'.format(epoch), unit='batch'):
            optimizer.zero_grad()
            if timeType == TimeType.ContinuousTimeLoss:
                loss = bfn.process_infinity(X.to(device))
            elif timeType == TimeType.DiscreteTimeLoss:
                loss = bfn.process_discrete(X.to(device), max_step=args.max_step)
            else:
                raise ValueError("The TimeType must be either ContinuousTimeLoss or DiscreteTimeLoss.")

            loss.backward()
            torch.nn.utils.clip_grad_norm_(unet.parameters(), 1)
            
            optimizer.step()
            ema.update()
            losses.append(loss.item())
            writer.add_scalar('Loss/train', loss.item(), num)
            num += 1
        scheduler.step()
        
        ave = 0
        for loss in losses:
            ave += loss
        ave = ave / len(losses)
        tqdm.write('Epoch {}: Loss: {:.8f}'.format(epoch, ave))
        writer.add_scalar('Loss/epoch_train', ave, epoch)
        if epoch % args.save_every_n_epoch == 0:
            torch.save(unet.state_dict(), f'./models/model_{epoch}.pth')
            tqdm.write('Epoch {}: saved: {}'.format(epoch, f'./models/model_{epoch}.pth'))

    torch.save(unet.state_dict(), args.save_model_path)
    tqdm.write('Epoch {}: saved: {}'.format(epoch, args.save_model_path))

def setup_train_common_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser.add_argument("--load_model_path", type=str, default=None)
    parser.add_argument("--save_model_path", type=str, default="./models/model.pth")
    parser.add_argument("--save_every_n_epoch", type=int, default=10)
    parser.add_argument("--data_path", type=str, default="./animeface")
    parser.add_argument("--batch", type=int, default=32)
    parser.add_argument("--epoch", type=int, default=100)
    parser.add_argument("--sigma", type=float, default=0.001)
    parser.add_argument("--height", type=int, default=32)
    parser.add_argument("--width", type=int, default=32)
    parser.add_argument("--lr", type=float, default=1e-4)
    # Discrete Time Loss Option
    parser.add_argument("--K", type=int, default=16)
    parser.add_argument("--max_step", type=int, default=1000)
    return parser


train_common.py
yoimers/bayesian-flow-networks-animeface
import argparse
from train_common import BFNType, TimeType, setup_train_common_parser, train

bfnType=BFNType.Continuous
timeType=TimeType.ContinuousTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_train_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    train(args, bfnType=bfnType, timeType=timeType)
train_continuous.py
yoimers/bayesian-flow-networks-animeface
import argparse
from train_common import BFNType, TimeType, setup_train_common_parser, train

bfnType=BFNType.Continuous
timeType=TimeType.DiscreteTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_train_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    train(args, bfnType=bfnType, timeType=timeType)
train_continuous_discrete_loss.py
yoimers/bayesian-flow-networks-animeface
import argparse
from train_common import BFNType, TimeType, setup_train_common_parser, train

bfnType=BFNType.Discretized
timeType=TimeType.ContinuousTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_train_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    train(args, bfnType=bfnType, timeType=timeType)
train_discretized.py
yoimers/bayesian-flow-networks-animeface
import argparse
from train_common import BFNType, TimeType, setup_train_common_parser, train

bfnType=BFNType.Discretized
timeType=TimeType.DiscreteTimeLoss

def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser = setup_train_common_parser(parser)
    return parser
    
if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    train(args, bfnType=bfnType, timeType=timeType)
train_discretized_discrete_loss.py
yoimers/bayesian-flow-networks-animeface
import random
from typing import Optional
from PIL import Image
import torchvision.transforms as transforms
import torch


def resize(h=32, w=32):
    return lambda img: img.resize((h, w), Image.BICUBIC)

def convert_to_3_channels(img):
    return img.convert('RGB')

def default_transform(h=32, w=32):
    return transforms.Compose([
        resize(h=h, w=w),
        convert_to_3_channels,
        transforms.ToTensor(),  # Convert images to tensors
        # transforms.RandomHorizontalFlip(p=0.5),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize the pixel values to [-1, 1]
    ])
    
    
def set_seed(seed: Optional[int]):
    if seed is None:
        seed = random.randint(1, 10000)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
utils.py
yoimers/bayesian-flow-networks-animeface