FluxTransformer2DModel(
(pos_embed): FluxPosEmbed()
(time_text_embed): CombinedTimestepGuidanceTextProjEmbeddings(
(time_proj): Timesteps()
(timestep_embedder): TimestepEmbedding(
(linear_1): Linear(in_features=256, out_features=3072, bias=True)
(act): SiLU()
(linear_2): Linear(in_features=3072, out_features=3072, bias=True)
)
(guidance_embedder): TimestepEmbedding(
(linear_1): Linear(in_features=256, out_features=3072, bias=True)
(act): SiLU()
(linear_2): Linear(in_features=3072, out_features=3072, bias=True)
)
(text_embedder): PixArtAlphaTextProjection(
(linear_1): Linear(in_features=768, out_features=3072, bias=True)
(act_1): SiLU()
(linear_2): Linear(in_features=3072, out_features=3072, bias=True)
)
)
(context_embedder): Linear(in_features=4096, out_features=3072, bias=True)
(x_embedder): Linear(in_features=64, out_features=3072, bias=True)
(transformer_blocks): ModuleList(
(0-18): 19 x FluxTransformerBlock(
(norm1): AdaLayerNormZero(
(silu): SiLU()
(linear): Linear(in_features=3072, out_features=18432, bias=True)
(norm): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
)
(norm1_context): AdaLayerNormZero(
(silu): SiLU()
(linear): Linear(in_features=3072, out_features=18432, bias=True)
(norm): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
)
(attn): FluxAttention(
(norm_q): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(norm_k): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(to_q): Linear(in_features=3072, out_features=3072, bias=True)
(to_k): Linear(in_features=3072, out_features=3072, bias=True)
(to_v): Linear(in_features=3072, out_features=3072, bias=True)
(to_out): ModuleList(
(0): Linear(in_features=3072, out_features=3072, bias=True)
(1): Dropout(p=0.0, inplace=False)
)
(norm_added_q): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(norm_added_k): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(add_q_proj): Linear(in_features=3072, out_features=3072, bias=True)
(add_k_proj): Linear(in_features=3072, out_features=3072, bias=True)
(add_v_proj): Linear(in_features=3072, out_features=3072, bias=True)
(to_add_out): Linear(in_features=3072, out_features=3072, bias=True)
)
(norm2): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
(ff): FeedForward(
(net): ModuleList(
(0): GELU(
(proj): Linear(in_features=3072, out_features=12288, bias=True)
)
(1): Dropout(p=0.0, inplace=False)
(2): Linear(in_features=12288, out_features=3072, bias=True)
)
)
(norm2_context): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
(ff_context): FeedForward(
(net): ModuleList(
(0): GELU(
(proj): Linear(in_features=3072, out_features=12288, bias=True)
)
(1): Dropout(p=0.0, inplace=False)
(2): Linear(in_features=12288, out_features=3072, bias=True)
)
)
)
)
(single_transformer_blocks): ModuleList(
(0-37): 38 x FluxSingleTransformerBlock(
(norm): AdaLayerNormZeroSingle(
(silu): SiLU()
(linear): Linear(in_features=3072, out_features=9216, bias=True)
(norm): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
)
(proj_mlp): Linear(in_features=3072, out_features=12288, bias=True)
(act_mlp): GELU(approximate='tanh')
(proj_out): Linear(in_features=15360, out_features=3072, bias=True)
(attn): FluxAttention(
(norm_q): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(norm_k): RMSNorm((128,), eps=1e-06, elementwise_affine=True)
(to_q): Linear(in_features=3072, out_features=3072, bias=True)
(to_k): Linear(in_features=3072, out_features=3072, bias=True)
(to_v): Linear(in_features=3072, out_features=3072, bias=True)
)
)
)
(norm_out): AdaLayerNormContinuous(
(silu): SiLU()
(linear): Linear(in_features=3072, out_features=6144, bias=True)
(norm): LayerNorm((3072,), eps=1e-06, elementwise_affine=False)
)
(proj_out): Linear(in_features=3072, out_features=64, bias=True)
)