credit.models.unet_diffusion

credit.models.unet_diffusion#

Attributes#

unet_config

Classes#

`RMSNorm`	Base class for all neural network modules.
`SinusoidalPosEmb`	Base class for all neural network modules.
`RandomOrLearnedSinusoidalPosEmb`	following @crowsonkb 's lead with random (learned optional) sinusoidal pos emb
`Block`	Base class for all neural network modules.
`ResnetBlock`	Base class for all neural network modules.
`LinearAttention`	Base class for all neural network modules.
`Attention`	Base class for all neural network modules.
`PeriodicConv2d`	Base class for all neural network modules.
`UnetDiffusion`	Base class for all neural network modules.

Functions#

`Upsample`(dim[, dim_out])
`Downsample`(dim[, dim_out])
`create_model`(config[, self_condition])	Initialize and return the CrossFormer model using a config dictionary.
`create_diffusion`(model, config)	Initialize and return the Gaussian Diffusion process.

Module Contents#

credit.models.unet_diffusion.Upsample(dim, dim_out=None)#

credit.models.unet_diffusion.Downsample(dim, dim_out=None)#

class credit.models.unet_diffusion.RMSNorm(dim)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

scale#

g#

forward(x)#

class credit.models.unet_diffusion.SinusoidalPosEmb(dim: int, theta: float = 10000)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

dim#

theta = 10000#

forward(x)#

class credit.models.unet_diffusion.RandomOrLearnedSinusoidalPosEmb(dim, is_random=False)#

Bases: torch.nn.Module

following @crowsonkb ‘s lead with random (learned optional) sinusoidal pos emb

weights#

forward(x)#

class credit.models.unet_diffusion.Block(dim, dim_out, dropout=0.0)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

proj#

norm#

act#

dropout#

forward(x, scale_shift=None)#

class credit.models.unet_diffusion.ResnetBlock(dim, dim_out, *, time_emb_dim=None, dropout=0.0)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

mlp#

block1#

block2#

res_conv#

forward(x, time_emb=None)#

class credit.models.unet_diffusion.LinearAttention(dim, heads=4, dim_head=32, num_mem_kv=4)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

scale = 0.1767766952966369#

heads = 4#

norm#

mem_kv#

to_qkv#

to_out#

forward(x)#

class credit.models.unet_diffusion.Attention(dim, heads=4, dim_head=32, num_mem_kv=4, flash=False)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

heads = 4#

norm#

attend#

mem_kv#

to_qkv#

to_out#

forward(x)#

class credit.models.unet_diffusion.PeriodicConv2d(dim_in, dim_out, kernel_size, padding=1)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

padding = 1#

conv#

forward(x)#

class credit.models.unet_diffusion.UnetDiffusion(image_height: int = 640, image_width: int = 1280, init_dim=None, frames: int = 2, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, dim: tuple = (64, 128, 256, 512), depth: tuple = (2, 2, 8, 2), dim_head: int = 32, padding_conf: dict = None, post_conf: dict = None, dim_mults: tuple = (1, 2, 4, 8), conditional_dimensions: int = 0, learned_variance: bool = False, learned_sinusoidal_cond: bool = False, random_fourier_features: bool = False, learned_sinusoidal_dim: int = 16, sinusoidal_pos_emb_theta: int = 10000, dropout: float = 0.0, attn_dim_head: int = 32, attn_heads: int = 4, full_attn: dict = None, flash_attn: bool = False, self_condition: bool = False, condition: bool = False, *args, **kwargs)#

Bases: credit.models.base_model.BaseModel

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

output_channels = 67#

channels = 4#

channels_out = 67#

pre_out_dim = 67#

self_condition = False#

conditional_dimensions = 0#

condition = False#

image_height = 640#

image_width = 1280#

frames = 2#

surface_channels = 7#

levels = 15#

use_post_block#

use_padding#

input_only_channels = 3#

input_channels = 70#

random_or_learned_sinusoidal_cond = False#

time_mlp#

downs#

ups#

mid_block1#

mid_attn#

mid_block2#

final_res_block#

final_conv#

property downsample_factor#

forward(x, time, x_self_cond=None, x_cond=None)#

credit.models.unet_diffusion.create_model(config, self_condition=True)#: Initialize and return the CrossFormer model using a config dictionary.

credit.models.unet_diffusion.create_diffusion(model, config)#: Initialize and return the Gaussian Diffusion process.

credit.models.unet_diffusion.unet_config#