fix: parameters order assumption in diffuser

[kfachikov]

Description

This change ensures that the inference handler is compatible with diffuser pipelines that do not have "prompt" as a first parameter, such as Flux2KleinPipeline.

Related Issue

Many diffusers pipelines, such as the Qwen one, have "prompt" as a first parameter. The existing code relied on this assumption - as we can see here,

The first element of the batch is passed as input to the model.

However, not all diffusers pipelines have the same signature. Flux2KleinPipeline for example expects model on first position and prompt on second. Running a pruna evaluation with the Flux2 model resulted in an error

ValueError: Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.

Important

This change does not remove the resolve the following hypothetical problem.

We have a dataset and its first batch element is a dict that is meant to be passed as a single value to one argument (rather than spread). Neither the old main path nor the new code does that — the whole framework already assumes that a dict should spread as kwargs.

How?

The change follows the already established patter to unwrap the prepared inputs - the way it's done in PrunaModel.run_inference used in evaluation_agent, the elapsed time metric, and handlethe memory metric.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor (no functional change)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

Testing

• I added or updated tests covering my changes
• Existing tests pass locally (uv run pytest -m "cpu and not slow")

For full setup and testing instructions, see the Contributing Guide.

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my code, especially for agent-assisted changes
• I updated the documentation where necessary

---

Thanks for contributing to Pruna! We're excited to review your work.

New to contributing? Check out our Contributing Guide for everything you need to get started.

Note:

• Draft PRs or PRs without a clear and detailed overview may be delayed.
• Please mark your PR as Ready for Review and ensure the sections above are filled out.
• Contributions that are entirely AI-generated without meaningful human review are discouraged.

[kfachikov]

@begumcig What do you think about reusing run_inference in the metric_energy and the metric_model_architecture`?

I see that you have made such a change in some other metrics a year ago (e.g., here) and wondered whether there was a particular reason why these two were not changed.

One small concern with metric_energy is that the current code avoids constantly moving the batch to the "device" and preparing the inputs (as it does it only once - code here. In case we use run_inference, this will happen on every iteration, possibly skewing the measurements.

AFAIU, the metric_model_architecture would not be affected.