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.

Extra

• Updated the documentation for DiffuserHandler, specifying the expected behavior for prepare_inputs.
• Updated the documentation for the collate.py module, elaborating on the return types for each collate method.

Related Issue

Many diffusers pipelines, such as QwenImagePipeline, 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 is backward compatible and does not break any existing working scenarios. It expands the usecases to diffusion models which do not expect prompt on first position (e.g., Flux2KleinPipeline), but does NOT resolve the following hypothetical problems/limitations.

• 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).
• Any diffusion model wrapped by PrunaModel and run with run_inference would expect the model inputs to be the first element in the batch passed as argument.

Neither the old implementation nor the new code handle these.

How?

The change in the metrics 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)
• Documentation update

Testing

• I added or updated tests covering my changes
  • I added comprehensive tests for the handler_diffusers module. They exercise the prepare_inputs behavior under
    • pruna native datasets (i.e., results from collate - image_generation_collate, prompt_collate, and prompt_with_auxiliaries_collate),
    • and custom datasets (i.e., including image for image-edit, or negative prompt for conditional generation)
• 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

TL;DR

I explored all the pruna native datasets (resolvable through PrunaDataModule.fromString) and the way they are processed (i.e., the collate functions). Here is a list with their groups by collate:

1. image_generation_collate
2. prompt_collate
3. prompt_with_auxiliaries_collate
4. audio_collate
5. image_classification_collate
6. text_generation_collate
7. question_answering_collate

Considering our scope - diffusion models, only Image Generation tasks are relevant. This limits the datasets to 1, 2, and 3. The audio collate is not appropriate for audio generation as it provides only the original audio and its transcription, without any description of what's in the audio.

Furthermore, diffusion models can also be used for the other tasks (image classification, text generation, or Q&A), even though they are not really popular and require careful manipulation. The structure provided by 5, 6, and 7 does not match the desired custom structure requested by these models and thus we do not support them.

---

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.

[kfachikov]

@begumcig

An update on using run_inference in the two metrics - I still have to run experiments and evaluate whether a (or any) difference is observed. As this change is just a QoL, I will leave it out of this PR.

edit

Posting experiment numbers here for history.

[gsprochette]

Looks good, @begumcig knows more than me about the inference handler so maybe she'll find missing edge cases but I didn't see one :) Thanks a lot !