Skip to content

Commit 9884ed2

Browse files
authored
FLUX.2 small decoder (#13428)
Add optional decoder_block_out_channels parameter to AutoencoderKLFlux2
1 parent 039e688 commit 9884ed2

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/diffusers/models/autoencoders/autoencoder_kl_flux2.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ def __init__(
9191
512,
9292
512,
9393
),
94+
decoder_block_out_channels: tuple[int, ...] | None = None,
9495
layers_per_block: int = 2,
9596
act_fn: str = "silu",
9697
latent_channels: int = 32,
@@ -124,7 +125,7 @@ def __init__(
124125
in_channels=latent_channels,
125126
out_channels=out_channels,
126127
up_block_types=up_block_types,
127-
block_out_channels=block_out_channels,
128+
block_out_channels=decoder_block_out_channels or block_out_channels,
128129
layers_per_block=layers_per_block,
129130
norm_num_groups=norm_num_groups,
130131
act_fn=act_fn,

0 commit comments

Comments
 (0)