Class StableCascadeUNet

src/diffusers/models/unets/unet_stable_cascade.py:136–583

class StableCascadeUNet(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        in_channels: int = 16,
        out_channels: int = 16,
        timestep_ratio_embedding_dim: int = 64,
        patch_size: int = 1,
        conditioning_dim: int = 2048,
        block_out_channels: tuple[int, ...] = (2048, 2048),
        num_attention_heads: tuple[int, ...] = (32, 32),
        down_num_layers_per_block: tuple[int, ...] = (8, 24),
        up_num_layers_per_block: tuple[int, ...] = (24, 8),
        down_blocks_repeat_mappers: tuple[int] | None = (
            1,
            1,
        ),
        up_blocks_repeat_mappers: tuple[int] | None = (1, 1),
        block_types_per_layer: tuple[tuple[str]] = (
            ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"),
            ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"),
        ),
        clip_text_in_channels: int | None = None,
        clip_text_pooled_in_channels=1280,
        clip_image_in_channels: int | None = None,
        clip_seq=4,
        effnet_in_channels: int | None = None,
        pixel_mapper_in_channels: int | None = None,
        kernel_size=3,
        dropout: float | tuple[float] = (0.1, 0.1),
        self_attn: bool | tuple[bool] = True,
        timestep_conditioning_type: tuple[str, ...] = ("sca", "crp"),
        switch_level: tuple[bool] | None = None,
    ):
        """

        Parameters:
            in_channels (`int`, defaults to 16):
                Number of channels in the input sample.
            out_channels (`int`, defaults to 16):
                Number of channels in the output sample.
            timestep_ratio_embedding_dim (`int`, defaults to 64):
                Dimension of the projected time embedding.
            patch_size (`int`, defaults to 1):
                Patch size to use for pixel unshuffling layer
            conditioning_dim (`int`, defaults to 2048):
                Dimension of the image and text conditional embedding.
            block_out_channels (tuple[int], defaults to (2048, 2048)):
                tuple of output channels for each block.
            num_attention_heads (tuple[int], defaults to (32, 32)):
                Number of attention heads in each attention block. Set to -1 to if block types in a layer do not have
                attention.
            down_num_layers_per_block (tuple[int], defaults to [8, 24]):
                Number of layers in each down block.
            up_num_layers_per_block (tuple[int], defaults to [24, 8]):
                Number of layers in each up block.