credit.models.wxformer.crossformer_diffusion

credit.models.wxformer.crossformer_diffusion#

Attributes#

`logger`
`crossformer_config`

Classes#

`ModelPrediction`
`UpBlock`	Base class for all neural network modules.
`Transformer`	Base class for all neural network modules.
`CrossFormerDiffusion`	Base class for all neural network modules.

Functions#

`exists`(x)
`default`(val, d)
`normalize_to_neg_one_to_one`(img)
`extract`(a, t, x_shape)
`identity`(t, args, *kwargs)
`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.wxformer.crossformer_diffusion.logger#

credit.models.wxformer.crossformer_diffusion.exists(x)#

credit.models.wxformer.crossformer_diffusion.default(val, d)#

credit.models.wxformer.crossformer_diffusion.normalize_to_neg_one_to_one(img)#

credit.models.wxformer.crossformer_diffusion.extract(a, t, x_shape)#

credit.models.wxformer.crossformer_diffusion.identity(t, *args, **kwargs)#

class credit.models.wxformer.crossformer_diffusion.ModelPrediction#

Bases: tuple

pred_noise#

pred_x_start#

class credit.models.wxformer.crossformer_diffusion.UpBlock(in_chans, out_chans, num_groups, num_residuals=2, emb_dim=None)#

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.

conv#

output_channels#

b#

to_emb = None#

forward(x, emb=None)#

class credit.models.wxformer.crossformer_diffusion.Transformer(dim, *, local_window_size, global_window_size, depth=4, dim_head=32, attn_dropout=0.0, ff_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.

layers#

forward(x, t_emb=None)#

class credit.models.wxformer.crossformer_diffusion.CrossFormerDiffusion(self_condition: bool = False, condition: bool = True, image_height: int = 640, patch_height: int = 1, image_width: int = 1280, patch_width: int = 1, 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, global_window_size: tuple = (5, 5, 2, 1), local_window_size: int = 10, cross_embed_kernel_sizes: tuple = ((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), cross_embed_strides: tuple = (4, 2, 2, 2), attn_dropout: float = 0.0, ff_dropout: float = 0.0, use_spectral_norm: bool = True, interp: bool = True, padding_conf: dict = None, post_conf: dict = None, **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.

image_height = 640#

image_width = 1280#

patch_height = 1#

patch_width = 1#

frames = 2#

channels = 4#

surface_channels = 7#

levels = 15#

use_spectral_norm = True#

use_interp = True#

use_padding#

use_post_block#

input_only_channels = 3#

input_channels = 70#

output_channels = 67#

total_input_channels = 70#

condition = True#

self_condition = False#

layers#

cube_embedding#

time_to_emb#

time_emb_proj#

up_block1#

up_block2#

up_block3#

up_block4#

forward(x, t, x_self_cond=False, x_cond=None)#

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

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

credit.models.wxformer.crossformer_diffusion.crossformer_config#