(
self,
sample_size: int = 65536,
sample_rate: int | None = None,
in_channels: int = 2,
out_channels: int = 2,
extra_in_channels: int = 0,
time_embedding_type: str = "fourier",
time_embedding_dim: int | None = None,
flip_sin_to_cos: bool = True,
use_timestep_embedding: bool = False,
freq_shift: float = 0.0,
down_block_types: tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"),
up_block_types: tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"),
mid_block_type: str = "UNetMidBlock1D",
out_block_type: str = None,
block_out_channels: tuple[int, ...] = (32, 32, 64),
act_fn: str = None,
norm_num_groups: int = 8,
layers_per_block: int = 1,
downsample_each_block: bool = False,
)
| 74 | |
| 75 | @register_to_config |
| 76 | def __init__( |
| 77 | self, |
| 78 | sample_size: int = 65536, |
| 79 | sample_rate: int | None = None, |
| 80 | in_channels: int = 2, |
| 81 | out_channels: int = 2, |
| 82 | extra_in_channels: int = 0, |
| 83 | time_embedding_type: str = "fourier", |
| 84 | time_embedding_dim: int | None = None, |
| 85 | flip_sin_to_cos: bool = True, |
| 86 | use_timestep_embedding: bool = False, |
| 87 | freq_shift: float = 0.0, |
| 88 | down_block_types: tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"), |
| 89 | up_block_types: tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"), |
| 90 | mid_block_type: str = "UNetMidBlock1D", |
| 91 | out_block_type: str = None, |
| 92 | block_out_channels: tuple[int, ...] = (32, 32, 64), |
| 93 | act_fn: str = None, |
| 94 | norm_num_groups: int = 8, |
| 95 | layers_per_block: int = 1, |
| 96 | downsample_each_block: bool = False, |
| 97 | ): |
| 98 | super().__init__() |
| 99 | self.sample_size = sample_size |
| 100 | |
| 101 | # time |
| 102 | if time_embedding_type == "fourier": |
| 103 | time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 |
| 104 | if time_embed_dim % 2 != 0: |
| 105 | raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") |
| 106 | self.time_proj = GaussianFourierProjection( |
| 107 | embedding_size=time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos |
| 108 | ) |
| 109 | timestep_input_dim = time_embed_dim |
| 110 | elif time_embedding_type == "positional": |
| 111 | time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 |
| 112 | self.time_proj = Timesteps( |
| 113 | block_out_channels[0], flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=freq_shift |
| 114 | ) |
| 115 | timestep_input_dim = block_out_channels[0] |
| 116 | else: |
| 117 | raise ValueError( |
| 118 | f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." |
| 119 | ) |
| 120 | |
| 121 | if use_timestep_embedding: |
| 122 | time_embed_dim = block_out_channels[0] * 4 |
| 123 | self.time_mlp = TimestepEmbedding( |
| 124 | in_channels=timestep_input_dim, |
| 125 | time_embed_dim=time_embed_dim, |
| 126 | act_fn=act_fn, |
| 127 | out_dim=block_out_channels[0], |
| 128 | ) |
| 129 | |
| 130 | self.down_blocks = nn.ModuleList([]) |
| 131 | self.mid_block = None |
| 132 | self.up_blocks = nn.ModuleList([]) |
| 133 | self.out_block = None |
nothing calls this directly
no test coverage detected