[OpenVINO] Support Flux.2-Klein Models on OpenVINO

[osabnis]

What does this PR do?

This PR enables OpenVINO export for black-forest-labs/FLUX.2-klein-4B by making the Flux export path more robust to Flux2/Klein-specific config and signature differences.

Main changes:

1. optimum/exporters/openvino/convert.py
   • Added adaptive Flux text encoder model-type resolution (clip-text vs gemma2-text-encoder vs t5-encoder-model).
   • Added handling for Flux pipelines where text encoders are wrapped as *ForCausalLM (exports backbone .model when applicable).
   • Added guarded pooled_projection_dim setup only when the transformer forward signature expects pooled_projections.
2. optimum/exporters/openvino/model_configs.py
   • Added dynamic Flux RoPE ID width inference helper (_get_flux_ids_dim).
   • Updated Flux dummy input generators to create img_ids / txt_ids using inferred ID width (instead of hardcoded 3).
   • Updated FluxTransformerOpenVINOConfig to include/exclude pooled_projections dummy/input mapping based on model config.

These changes keep the existing Flux path intact while supporting Flux2-Klein variants.

Example of conversion cmd-line for black-forest-labs/FLUX.2-klein-4B:

optimum-cli export openvino --model black-forest-labs/FLUX.2-klein-4B ov_flux.2-klein-4b

Example of inference for black-forest-labs/FLUX.2-klein-4B using OpenVINO backend:

from pathlib import Path
import numpy as np
import torch
from optimum.intel import OVDiffusionPipeline

model_dir = Path("./ov_flux.2-klein-4b")

# deterministic seeds for reproducibility
np.random.seed(0)

pipe = OVDiffusionPipeline.from_pretrained(model_dir, compile=False)
pipe.to("GPU")
pipe.compile()

out = pipe(
    prompt="Astronaut on the moon",
    height=512,
    width=512,
    num_inference_steps=4,
    output_type="pil",
)

img = out.images[0]
img.save("flux2_klein_smoke.png")
print("Saved:", Path("flux2_klein_smoke.png").resolve())
print("Image size:", img.size)

Converted Model Output from above script:

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[Copilot]

Pull request overview

This PR extends the Optimum-Intel OpenVINO export + inference path to handle Flux.2 “Klein” variants (e.g., black-forest-labs/FLUX.2-klein-4B) by making Flux component export/configuration more adaptive (text encoder type resolution, dynamic RoPE id width, optional pooled projections) and registering a new OV pipeline class for Flux2Klein.

Changes:

• Add Flux text-encoder type resolution and make Flux export more tolerant of *ForCausalLM-wrapped encoders and signature differences.
• Infer Flux RoPE id width dynamically and use it in Flux dummy input generation; conditionally include pooled projections.
• Introduce OVFlux2KleinPipeline and register it in the OpenVINO diffusion pipeline mappings.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
optimum/exporters/openvino/convert.py	Adds Flux text-encoder model-type resolution, pooled-projection guarding, and ONNX fallback in export_models.
optimum/exporters/openvino/model_configs.py	Adds _get_flux_ids_dim, updates Flux dummy generators to use inferred id width, and makes pooled projections optional for Flux transformer export config.
optimum/intel/openvino/modeling_diffusion.py	Adds Flux2Klein pipeline support and makes transformer reshape / transformer inputs more flexible for ids + hidden-states behavior.
optimum/intel/openvino/__init__.py	Re-exports OVFlux2KleinPipeline when diffusers is available.
optimum/intel/__init__.py	Adds OVFlux2KleinPipeline to lazy import structures and dummy objects mapping.
optimum/intel/utils/dummy_openvino_and_diffusers_objects.py	Adds dummy OVFlux2KleinPipeline for optional dependency environments.
optimum/exporters/openvino/__init__.py	Re-exports _resolve_flux_text_encoder_model_type from convert.py.
tests/openvino/test_export.py	Adds unit tests for _get_flux_ids_dim and _resolve_flux_text_encoder_model_type.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.