[Bugfix] Fix Flux2 Dev Guidance

[alex-jw-brooks]

Purpose

Flux2 dev is currently not passing guidance through correctly; regardless of the sampling params, the guidance tensor passed to the transformer is hardcoded to None.

This PR creates the guidance tensor and passes it through to the timestep guidance embeddings (original pipeline for reference here).

Test Result

To validate, you can run an example like this:

from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.entrypoints.omni import Omni

MODEL = "black-forest-labs/FLUX.2-dev"

PROMPT = "A white cat sitting on a black table near a windows with visible sunlight."
SEED = 42
NUM_STEPS = 4

def generate_with_guidance(omni: Omni, guidance_scale: float, output_path: str):
    """Generate image with specific guidance_scale using existing Omni instance."""
    outputs = omni.generate(
        prompts=PROMPT,
        sampling_params_list=OmniDiffusionSamplingParams(
            seed=SEED,
            height=1024,
            width=1024,
            num_inference_steps=NUM_STEPS,
            guidance_scale=guidance_scale
        )
    )

    image = outputs[0].request_output.images[0]
    image.save(output_path)
    return image


if __name__ == "__main__":
    omni = Omni(
        model=MODEL,
        dtype="bfloat16",
        enable_cpu_offload=True,
    )

    img_low = generate_with_guidance(omni, 1.0, "guidance_1.0.png")
    img_mid = generate_with_guidance(omni, 4.0, "guidance_4.0.png")

In the current code, the result will be the same for both since the guidance isn't used, but we should see a poor result for guidance 1.0 and a pretty good one for 4.0.

[lishunyang12]

Looks good, one nit.

[gcanlin]

Could you please paste the generated image before and after this PR?

[SamitHuang]

LGTM

[alex-jw-brooks]

@gcanlin sure, with the example code, no matter what value you pass for guidance, you'll always get the same image since it's unused.

With the fix, we should see the guidance affect the result, so giving something like guidance 1.0 should be much worse than guidance 4.0 for example, which is closer to what is recommended for this model.

(left is guidance 1.0, right is 4.0, with otherwise identical configurations)