r""" A 1D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Parameters: sam
| 39 | |
| 40 | |
| 41 | class UNet1DModel(ModelMixin, ConfigMixin): |
| 42 | r""" |
| 43 | A 1D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. |
| 44 | |
| 45 | This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented |
| 46 | for all models (such as downloading or saving). |
| 47 | |
| 48 | Parameters: |
| 49 | sample_size (`int`, *optional*): Default length of sample. Should be adaptable at runtime. |
| 50 | in_channels (`int`, *optional*, defaults to 2): Number of channels in the input sample. |
| 51 | out_channels (`int`, *optional*, defaults to 2): Number of channels in the output. |
| 52 | extra_in_channels (`int`, *optional*, defaults to 0): |
| 53 | Number of additional channels to be added to the input of the first down block. Useful for cases where the |
| 54 | input data has more channels than what the model was initially designed for. |
| 55 | time_embedding_type (`str`, *optional*, defaults to `"fourier"`): Type of time embedding to use. |
| 56 | freq_shift (`float`, *optional*, defaults to 0.0): Frequency shift for Fourier time embedding. |
| 57 | flip_sin_to_cos (`bool`, *optional*, defaults to `False`): |
| 58 | Whether to flip sin to cos for Fourier time embedding. |
| 59 | down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")`): |
| 60 | Tuple of downsample block types. |
| 61 | up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")`): |
| 62 | Tuple of upsample block types. |
| 63 | block_out_channels (`Tuple[int]`, *optional*, defaults to `(32, 32, 64)`): |
| 64 | Tuple of block output channels. |
| 65 | mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock1D"`): Block type for middle of UNet. |
| 66 | out_block_type (`str`, *optional*, defaults to `None`): Optional output processing block of UNet. |
| 67 | act_fn (`str`, *optional*, defaults to `None`): Optional activation function in UNet blocks. |
| 68 | norm_num_groups (`int`, *optional*, defaults to 8): The number of groups for normalization. |
| 69 | layers_per_block (`int`, *optional*, defaults to 1): The number of layers per block. |
| 70 | downsample_each_block (`int`, *optional*, defaults to `False`): |
| 71 | Experimental feature for using a UNet without upsampling. |
| 72 | """ |
| 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 |
no outgoing calls