vllm_omni.diffusion.models.flux2_klein ¶

attention_kwargs

current_timestep

default_sample_size = 128

do_classifier_free_guidance

guidance_scale

image_processor = Flux2ImageProcessor(
    vae_scale_factor=self.vae_scale_factor * 2
)

interrupt

is_distilled = is_distilled

latent_channels = (
    self.vae.config.latent_channels
    if hasattr(self.vae, "config")
    else 16
)

mask_processor = VaeImageProcessor(
    vae_scale_factor=self.vae_scale_factor * 2,
    vae_latent_channels=self.latent_channels,
    do_normalize=False,
    do_binarize=True,
    do_convert_grayscale=True,
)

num_timesteps

od_config = od_config

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    model,
    subfolder="scheduler",
    local_files_only=local_files_only,
)

support_image_input = True

text_encoder = from_pretrained_with_prefetch(
    Qwen3ForCausalLM.from_pretrained,
    model,
    subfolder="text_encoder",
    prefetch_list=flux2_subfolders,
    local_files_only=local_files_only,
).to(self._execution_device)

tokenizer = Qwen2TokenizerFast.from_pretrained(
    model,
    subfolder="tokenizer",
    local_files_only=local_files_only,
)

tokenizer_max_length = 512

transformer = Flux2Transformer2DModel(
    quant_config=od_config.quantization_config,
    **transformer_kwargs,
)

vae = from_pretrained_with_prefetch(
    AutoencoderKLFlux2.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=flux2_subfolders,
    local_files_only=local_files_only,
).to(self._execution_device)

vae_scale_factor = (
    2 ** (len(self.vae.config.block_out_channels) - 1)
    if getattr(self, "vae", None)
    else 8
)

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder="transformer",
        revision=None,
        prefix="transformer.",
        fall_back_to_pt=True,
    )
]

config = SimpleNamespace(
    patch_size=patch_size,
    in_channels=in_channels,
    out_channels=self.out_channels,
    num_layers=num_layers,
    num_single_layers=num_single_layers,
    attention_head_dim=attention_head_dim,
    num_attention_heads=num_attention_heads,
    joint_attention_dim=joint_attention_dim,
    timestep_guidance_channels=timestep_guidance_channels,
    mlp_ratio=mlp_ratio,
    axes_dims_rope=axes_dims_rope,
    rope_theta=rope_theta,
    eps=eps,
    guidance_embeds=guidance_embeds,
)

context_embedder = nn.Linear(
    joint_attention_dim, self.inner_dim, bias=False
)

double_stream_modulation_img = Flux2Modulation(
    self.inner_dim, mod_param_sets=2, bias=False
)

double_stream_modulation_txt = Flux2Modulation(
    self.inner_dim, mod_param_sets=2, bias=False
)

dtype: dtype

inner_dim = num_attention_heads * attention_head_dim

norm_out = AdaLayerNormContinuous(
    self.inner_dim,
    self.inner_dim,
    elementwise_affine=False,
    eps=eps,
    bias=False,
)

out_channels = out_channels or in_channels

parallel_config = od_config.parallel_config

pos_embed = Flux2PosEmbed(
    theta=rope_theta, axes_dim=list(axes_dims_rope)
)

proj_out = nn.Linear(
    self.inner_dim,
    patch_size * patch_size * self.out_channels,
    bias=False,
)

rope_prepare = Flux2RopePrepare(self.pos_embed)

single_stream_modulation = Flux2Modulation(
    self.inner_dim, mod_param_sets=1, bias=False
)

single_transformer_blocks = nn.ModuleList(
    [
        (
            Flux2SingleTransformerBlock(
                parallel_config=self.parallel_config,
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                mlp_ratio=mlp_ratio,
                eps=eps,
                bias=False,
                quant_config=quant_config,
                prefix=f"single_transformer_blocks.{i}",
            )
        )
        for i in (range(num_single_layers))
    ]
)

time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
    in_channels=timestep_guidance_channels,
    embedding_dim=self.inner_dim,
    bias=False,
    guidance_embeds=guidance_embeds,
)

transformer_blocks = nn.ModuleList(
    [
        (
            Flux2TransformerBlock(
                parallel_config=self.parallel_config,
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                mlp_ratio=mlp_ratio,
                eps=eps,
                bias=False,
                quant_config=quant_config,
                prefix=f"transformer_blocks.{i}",
            )
        )
        for i in (range(num_layers))
    ]
)

x_embedder = nn.Linear(
    in_channels, self.inner_dim, bias=False
)