(
self,
sample_size: int = 65536,
sample_rate: Optional[int] = None,
in_channels: int = 2,
out_channels: int = 2,
extra_in_channels: int = 0,
time_embedding_type: str = "fourier",
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: Tuple[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,
)
| 73 | |
| 74 | @register_to_config |
| 75 | def __init__( |
| 76 | self, |
| 77 | sample_size: int = 65536, |
| 78 | sample_rate: Optional[int] = None, |
| 79 | in_channels: int = 2, |
| 80 | out_channels: int = 2, |
| 81 | extra_in_channels: int = 0, |
| 82 | time_embedding_type: str = "fourier", |
| 83 | flip_sin_to_cos: bool = True, |
| 84 | use_timestep_embedding: bool = False, |
| 85 | freq_shift: float = 0.0, |
| 86 | down_block_types: Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"), |
| 87 | up_block_types: Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"), |
| 88 | mid_block_type: Tuple[str] = "UNetMidBlock1D", |
| 89 | out_block_type: str = None, |
| 90 | block_out_channels: Tuple[int] = (32, 32, 64), |
| 91 | act_fn: str = None, |
| 92 | norm_num_groups: int = 8, |
| 93 | layers_per_block: int = 1, |
| 94 | downsample_each_block: bool = False, |
| 95 | ): |
| 96 | super().__init__() |
| 97 | self.sample_size = sample_size |
| 98 | |
| 99 | # time |
| 100 | if time_embedding_type == "fourier": |
| 101 | self.time_proj = GaussianFourierProjection( |
| 102 | embedding_size=8, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos |
| 103 | ) |
| 104 | timestep_input_dim = 2 * block_out_channels[0] |
| 105 | elif time_embedding_type == "positional": |
| 106 | self.time_proj = Timesteps( |
| 107 | block_out_channels[0], flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=freq_shift |
| 108 | ) |
| 109 | timestep_input_dim = block_out_channels[0] |
| 110 | |
| 111 | if use_timestep_embedding: |
| 112 | time_embed_dim = block_out_channels[0] * 4 |
| 113 | self.time_mlp = TimestepEmbedding( |
| 114 | in_channels=timestep_input_dim, |
| 115 | time_embed_dim=time_embed_dim, |
| 116 | act_fn=act_fn, |
| 117 | out_dim=block_out_channels[0], |
| 118 | ) |
| 119 | |
| 120 | self.down_blocks = nn.ModuleList([]) |
| 121 | self.mid_block = None |
| 122 | self.up_blocks = nn.ModuleList([]) |
| 123 | self.out_block = None |
| 124 | |
| 125 | # down |
| 126 | output_channel = in_channels |
| 127 | for i, down_block_type in enumerate(down_block_types): |
| 128 | input_channel = output_channel |
| 129 | output_channel = block_out_channels[i] |
| 130 | |
| 131 | if i == 0: |
| 132 | input_channel += extra_in_channels |
nothing calls this directly
no test coverage detected