Skip to content

vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein

logger module-attribute 

logger = init_logger(__name__)

Flux2ImageProcessor

Bases: VaeImageProcessor

Image processor to preprocess the reference image for Flux2 klein.

check_image_input staticmethod 

check_image_input(
    image: Image,
    max_aspect_ratio: int = 8,
    min_side_length: int = 64,
    max_area: int = 1024 * 1024,
) -> Image

concatenate_images staticmethod 

concatenate_images(images: list[Image]) -> Image

Flux2KleinPipeline

Bases: Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery

Flux2 klein pipeline for text-to-image generation.

attention_kwargs property 

attention_kwargs

current_timestep property 

current_timestep

default_sample_size instance-attribute 

default_sample_size = 128

do_classifier_free_guidance property 

do_classifier_free_guidance

guidance_scale property 

guidance_scale

image_processor instance-attribute 

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

interrupt property 

interrupt

is_distilled instance-attribute 

is_distilled = is_distilled

latent_channels instance-attribute 

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

mask_processor instance-attribute 

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 property 

num_timesteps

od_config instance-attribute 

od_config = od_config

scheduler instance-attribute 

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

support_image_input class-attribute instance-attribute 

support_image_input = True

text_encoder instance-attribute 

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 instance-attribute 

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

tokenizer_max_length instance-attribute 

tokenizer_max_length = 512

transformer instance-attribute 

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

vae instance-attribute 

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 instance-attribute 

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

weights_sources instance-attribute 

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

check_inputs 

check_inputs(
    prompt,
    height,
    width,
    prompt_embeds=None,
    callback_on_step_end_tensor_inputs=None,
    guidance_scale=None,
    strength=None,
    num_inference_steps=None,
)

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    device: device | None = None,
    num_images_per_prompt: int = 1,
    prompt_embeds: Tensor | None = None,
    max_sequence_length: int = 512,
    text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
)

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

get_timesteps 

get_timesteps(num_inference_steps, strength, device)

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

prepare_image_latents 

prepare_image_latents(
    images: list[Tensor],
    batch_size,
    generator: Generator,
    device,
    dtype,
)

prepare_latents 

prepare_latents(
    batch_size,
    num_latents_channels,
    height,
    width,
    dtype,
    device,
    generator: Generator,
    latents: Tensor | None = None,
)

prepare_mask_latents 

prepare_mask_latents(
    mask,
    masked_image,
    batch_size,
    num_channels_latents,
    num_images_per_prompt,
    height,
    width,
    dtype,
    device,
    generator,
)

compute_empirical_mu 

compute_empirical_mu(
    image_seq_len: int, num_steps: int
) -> float

get_flux2_klein_post_process_func 

get_flux2_klein_post_process_func(
    od_config: OmniDiffusionConfig,
)