(
self,
in_channels: int = 4,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
down_block_types: Tuple[str] = (
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
only_cross_attention: Union[bool, Tuple[bool]] = False,
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
layers_per_block: int = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
act_fn: str = "silu",
norm_num_groups: Optional[int] = 32,
norm_eps: float = 1e-5,
cross_attention_dim: int = 1280,
attention_head_dim: Union[int, Tuple[int]] = 8,
use_linear_projection: bool = False,
class_embed_type: Optional[str] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
projection_class_embeddings_input_dim: Optional[int] = None,
controlnet_conditioning_channel_order: str = "rgb",
conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),
global_pool_conditions: bool = False,
)
| 91 | |
| 92 | @register_to_config |
| 93 | def __init__( |
| 94 | self, |
| 95 | in_channels: int = 4, |
| 96 | flip_sin_to_cos: bool = True, |
| 97 | freq_shift: int = 0, |
| 98 | down_block_types: Tuple[str] = ( |
| 99 | "CrossAttnDownBlock2D", |
| 100 | "CrossAttnDownBlock2D", |
| 101 | "CrossAttnDownBlock2D", |
| 102 | "DownBlock2D", |
| 103 | ), |
| 104 | only_cross_attention: Union[bool, Tuple[bool]] = False, |
| 105 | block_out_channels: Tuple[int] = (320, 640, 1280, 1280), |
| 106 | layers_per_block: int = 2, |
| 107 | downsample_padding: int = 1, |
| 108 | mid_block_scale_factor: float = 1, |
| 109 | act_fn: str = "silu", |
| 110 | norm_num_groups: Optional[int] = 32, |
| 111 | norm_eps: float = 1e-5, |
| 112 | cross_attention_dim: int = 1280, |
| 113 | attention_head_dim: Union[int, Tuple[int]] = 8, |
| 114 | use_linear_projection: bool = False, |
| 115 | class_embed_type: Optional[str] = None, |
| 116 | num_class_embeds: Optional[int] = None, |
| 117 | upcast_attention: bool = False, |
| 118 | resnet_time_scale_shift: str = "default", |
| 119 | projection_class_embeddings_input_dim: Optional[int] = None, |
| 120 | controlnet_conditioning_channel_order: str = "rgb", |
| 121 | conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256), |
| 122 | global_pool_conditions: bool = False, |
| 123 | ): |
| 124 | super().__init__() |
| 125 | |
| 126 | # Check inputs |
| 127 | if len(block_out_channels) != len(down_block_types): |
| 128 | raise ValueError( |
| 129 | f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." |
| 130 | ) |
| 131 | |
| 132 | if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): |
| 133 | raise ValueError( |
| 134 | f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." |
| 135 | ) |
| 136 | |
| 137 | if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): |
| 138 | raise ValueError( |
| 139 | f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." |
| 140 | ) |
| 141 | |
| 142 | # input |
| 143 | conv_in_kernel = 3 |
| 144 | conv_in_padding = (conv_in_kernel - 1) // 2 |
| 145 | self.conv_in = nn.Conv2d( |
| 146 | in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding |
| 147 | ) |
| 148 | |
| 149 | # time |
| 150 | time_embed_dim = block_out_channels[0] * 4 |
no test coverage detected