credit.models.unet_diffusion ============================ .. py:module:: credit.models.unet_diffusion Attributes ---------- .. autoapisummary:: credit.models.unet_diffusion.unet_config Classes ------- .. autoapisummary:: credit.models.unet_diffusion.RMSNorm credit.models.unet_diffusion.SinusoidalPosEmb credit.models.unet_diffusion.RandomOrLearnedSinusoidalPosEmb credit.models.unet_diffusion.Block credit.models.unet_diffusion.ResnetBlock credit.models.unet_diffusion.LinearAttention credit.models.unet_diffusion.Attention credit.models.unet_diffusion.PeriodicConv2d credit.models.unet_diffusion.UnetDiffusion Functions --------- .. autoapisummary:: credit.models.unet_diffusion.Upsample credit.models.unet_diffusion.Downsample credit.models.unet_diffusion.create_model credit.models.unet_diffusion.create_diffusion Module Contents --------------- .. py:function:: Upsample(dim, dim_out=None) .. py:function:: Downsample(dim, dim_out=None) .. py:class:: RMSNorm(dim) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: scale .. py:attribute:: g .. py:method:: forward(x) .. py:class:: SinusoidalPosEmb(dim: int, theta: float = 10000) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: dim .. py:attribute:: theta :value: 10000 .. py:method:: forward(x) .. py:class:: RandomOrLearnedSinusoidalPosEmb(dim, is_random=False) Bases: :py:obj:`torch.nn.Module` following @crowsonkb 's lead with random (learned optional) sinusoidal pos emb .. py:attribute:: weights .. py:method:: forward(x) .. py:class:: Block(dim, dim_out, dropout=0.0) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: proj .. py:attribute:: norm .. py:attribute:: act .. py:attribute:: dropout .. py:method:: forward(x, scale_shift=None) .. py:class:: ResnetBlock(dim, dim_out, *, time_emb_dim=None, dropout=0.0) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: mlp .. py:attribute:: block1 .. py:attribute:: block2 .. py:attribute:: res_conv .. py:method:: forward(x, time_emb=None) .. py:class:: LinearAttention(dim, heads=4, dim_head=32, num_mem_kv=4) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: scale :value: 0.1767766952966369 .. py:attribute:: heads :value: 4 .. py:attribute:: norm .. py:attribute:: mem_kv .. py:attribute:: to_qkv .. py:attribute:: to_out .. py:method:: forward(x) .. py:class:: Attention(dim, heads=4, dim_head=32, num_mem_kv=4, flash=False) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: heads :value: 4 .. py:attribute:: norm .. py:attribute:: attend .. py:attribute:: mem_kv .. py:attribute:: to_qkv .. py:attribute:: to_out .. py:method:: forward(x) .. py:class:: PeriodicConv2d(dim_in, dim_out, kernel_size, padding=1) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: padding :value: 1 .. py:attribute:: conv .. py:method:: forward(x) .. py:class:: 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: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: output_channels :value: 67 .. py:attribute:: channels :value: 4 .. py:attribute:: channels_out :value: 67 .. py:attribute:: pre_out_dim :value: 67 .. py:attribute:: self_condition :value: False .. py:attribute:: conditional_dimensions :value: 0 .. py:attribute:: condition :value: False .. py:attribute:: image_height :value: 640 .. py:attribute:: image_width :value: 1280 .. py:attribute:: frames :value: 2 .. py:attribute:: surface_channels :value: 7 .. py:attribute:: levels :value: 15 .. py:attribute:: use_post_block .. py:attribute:: use_padding .. py:attribute:: input_only_channels :value: 3 .. py:attribute:: input_channels :value: 70 .. py:attribute:: random_or_learned_sinusoidal_cond :value: False .. py:attribute:: time_mlp .. py:attribute:: downs .. py:attribute:: ups .. py:attribute:: mid_block1 .. py:attribute:: mid_attn .. py:attribute:: mid_block2 .. py:attribute:: final_res_block .. py:attribute:: final_conv .. py:property:: downsample_factor .. py:method:: forward(x, time, x_self_cond=None, x_cond=None) .. py:function:: create_model(config, self_condition=True) Initialize and return the CrossFormer model using a config dictionary. .. py:function:: create_diffusion(model, config) Initialize and return the Gaussian Diffusion process. .. py:data:: unet_config