From 04e1edf2a47dd4777adfc96819b3c2e401efc9ae Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Mon, 6 Jul 2026 14:00:35 +0000
Subject: [PATCH 1/6] fix: positional argumnets assumption in diffuser

This change ensures that the inference handler is compatible with
diffuser pipelines that do not have a leading positional `prompt`
argument, such as Flux2KleinPipeline.
---
 src/pruna/engine/handler/handler_diffuser.py           |  5 ++++-
 src/pruna/evaluation/metrics/metric_energy.py          | 10 ++++++++--
 .../evaluation/metrics/metric_model_architecture.py    |  5 ++++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/pruna/engine/handler/handler_diffuser.py b/src/pruna/engine/handler/handler_diffuser.py
index fd1fd239..a2597058 100644
--- a/src/pruna/engine/handler/handler_diffuser.py
+++ b/src/pruna/engine/handler/handler_diffuser.py
@@ -63,7 +63,10 @@ def prepare_inputs(
         Any
             The prepared inputs.
         """
-        if "prompt" in self.call_signature.parameters or "args" in self.call_signature.parameters:
+        if "prompt" in self.call_signature.parameters:
+            x, _ = batch
+            return x if isinstance(x, dict) else {"prompt": x}
+        elif "args" in self.call_signature.parameters:
             x, _ = batch
             return x
         else:  # Unconditional generation models
diff --git a/src/pruna/evaluation/metrics/metric_energy.py b/src/pruna/evaluation/metrics/metric_energy.py
index 8858f8de..7edf39a7 100644
--- a/src/pruna/evaluation/metrics/metric_energy.py
+++ b/src/pruna/evaluation/metrics/metric_energy.py
@@ -105,11 +105,17 @@ def compute(self, model: PrunaModel, dataloader: DataLoader) -> Dict[str, Any] |
 
         # Warmup
         for _ in tqdm(range(self.n_warmup_iterations), desc="Warm-up for energy consumption metric", unit="iter"):
-            model(inputs, **model.inference_handler.model_args)
+            if isinstance(inputs, dict):
+                model(**inputs, **model.inference_handler.model_args)
+            else:
+                model(inputs, **model.inference_handler.model_args)
 
         tracker.start_task("Inference")
         for _ in tqdm(range(self.n_iterations), desc="Measuring energy consumption", unit="iter"):
-            model(inputs, **model.inference_handler.model_args)
+            if isinstance(inputs, dict):
+                model(**inputs, **model.inference_handler.model_args)
+            else:
+                model(inputs, **model.inference_handler.model_args)
         tracker.stop_task()
 
         # Make sure all the operations are finished before stopping the tracker
diff --git a/src/pruna/evaluation/metrics/metric_model_architecture.py b/src/pruna/evaluation/metrics/metric_model_architecture.py
index 9afb05fe..af2ac4bc 100644
--- a/src/pruna/evaluation/metrics/metric_model_architecture.py
+++ b/src/pruna/evaluation/metrics/metric_model_architecture.py
@@ -88,7 +88,10 @@ def compute(self, model: PrunaModel, dataloader: DataLoader) -> Dict[str, Any] |
         batch = model.inference_handler.move_inputs_to_device(batch, self.device)
         inputs = model.inference_handler.prepare_inputs(batch)
 
-        model(inputs, **model.inference_handler.model_args)
+        if isinstance(inputs, dict):
+            model(**inputs, **model.inference_handler.model_args)
+        else:
+            model(inputs, **model.inference_handler.model_args)
 
         total_macs = 0
         self.module_macs = {}

From ab46faceddf60b031040eec20c4e67c86dca67ad Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Thu, 9 Jul 2026 08:04:15 +0000
Subject: [PATCH 2/6] docs: improve return type documentation in collate

---
 src/pruna/data/collate.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/pruna/data/collate.py b/src/pruna/data/collate.py
index 2803ee50..fadb1407 100644
--- a/src/pruna/data/collate.py
+++ b/src/pruna/data/collate.py
@@ -91,8 +91,9 @@ def image_generation_collate(
 
     Returns
     -------
-    Tuple[torch.Tensor, Any]
-        The collated data with size img_size and normalized to [0, 1].
+    Tuple[List[str], Union[Float[torch.Tensor, ImageShape], Int[torch.Tensor, ImageShape]]]
+        A tuple with all text prompts as first element, and all images as second element.
+        The images are resized to img_size and converted to the desired format.
     """
     transformations = image_format_to_transforms(output_format, img_size)
     image_col = _resolve_column(column_map, "image")
@@ -127,7 +128,7 @@ def prompt_collate(data: Any, column_map: dict[str, str] | None = None) -> Tuple
     Returns
     -------
     Tuple[List[str], None]
-        The collated data.
+        A tuple with all text prompts as first element, and None as second element.
     """
     text_col = _resolve_column(column_map, "text")
     return [item[text_col] for item in data], None
@@ -154,7 +155,7 @@ def prompt_with_auxiliaries_collate(
     Returns
     -------
     Tuple[List[str], Any]
-        The collated data.
+        A tuple with all text prompts as first element, and all auxiliary dictionaries as second element.
     """
     text_col = _resolve_column(column_map, "text")
     #  The text column has the prompt.
@@ -180,8 +181,8 @@ def audio_collate(data: Any, column_map: dict[str, str] | None = None) -> Tuple[
 
     Returns
     -------
-    List[str]
-        The collated data.
+    Tuple[List[str], List[str]]
+        A tuple with all audio paths as first element, and all text transcriptions as second element.
     """
     audio_col = _resolve_column(column_map, "audio")
     sentence_col = _resolve_column(column_map, "sentence")
@@ -209,8 +210,10 @@ def image_classification_collate(
 
     Returns
     -------
-    Tuple[torch.Tensor, torch.Tensor]
-        The collated data.
+    Tuple[Float[torch.Tensor, ImageShape], Int[torch.Tensor, LabelShape]]
+        A tuple with all images, stacked into a single tensor, as first element,
+        and all labels, stacked into a single tensor, as second element.
+        The images are resized to img_size and converted to the desired format.
     """
     transformations = image_format_to_transforms(output_format, img_size)
     image_col = _resolve_column(column_map, "image")
@@ -251,7 +254,8 @@ def text_generation_collate(
     Returns
     -------
     Tuple[torch.Tensor, torch.Tensor]
-        The collated data.
+        A tuple with tokens from all text prompts as first element,
+        and their corresponding next tokens as second element. 
     """
     text_col = _resolve_column(column_map, "text")
     input_ids = []
@@ -290,7 +294,8 @@ def question_answering_collate(
     Returns
     -------
     Tuple[torch.Tensor, torch.Tensor]
-        The collated data.
+        A tuple with tokenized questions as first element,
+        and tokenized answers as second element.
     """
     question_col = _resolve_column(column_map, "question")
     answer_col = _resolve_column(column_map, "answer")

From 3c92b12261c14ee8adbb9a5b1ad1c6b2e5905ef6 Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Thu, 9 Jul 2026 08:31:37 +0000
Subject: [PATCH 3/6] docs: improve documentation in handler_diffuser

---
 src/pruna/engine/handler/handler_diffuser.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/pruna/engine/handler/handler_diffuser.py b/src/pruna/engine/handler/handler_diffuser.py
index a2597058..33aaaf61 100644
--- a/src/pruna/engine/handler/handler_diffuser.py
+++ b/src/pruna/engine/handler/handler_diffuser.py
@@ -48,14 +48,21 @@ def __init__(self, call_signature: inspect.Signature, model_args: Optional[Dict[
         self.model_args = default_args
 
     def prepare_inputs(
-        self, batch: List[str] | torch.Tensor | Tuple[List[str] | torch.Tensor | dict[str, Any], ...] | dict[str, Any]
+        self, batch: Tuple[List[str] | torch.Tensor | dict[str, Any], ...]
     ) -> Any:
         """
         Prepare the inputs for the model.
 
+        The batch's first element is considered the model input.
+        1. The native pruna data format for generation is a tuple with the first element being a list with the prompts.
+        Check ``pruna/src/pruna/data/collate.py`` for more details.
+        2. To provide additional arguments to the model (e.g., negative prompts), \
+        construct the first element as a dictionary (e.g., ``{"prompt": [...], "negative_prompt": [...]}``).
+
+
         Parameters
         ----------
-        batch : List[str] | torch.Tensor | Tuple[List[str] | torch.Tensor | dict[str, Any], ...] | dict[str, Any]
+        batch : Tuple[List[str] | torch.Tensor | dict[str, Any], ...]
             The batch to prepare the inputs for.
 
         Returns

From e39735f0168a5a9fa03c6e5d26245fc1ba7c2842 Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Thu, 9 Jul 2026 08:54:05 +0000
Subject: [PATCH 4/6] test: add comprehensive test for handler_diffuser

---
 tests/engine/handler/test_handler_diffuser.py | 241 ++++++++++++++++++
 tests/fixtures.py                             |   1 +
 2 files changed, 242 insertions(+)
 create mode 100644 tests/engine/handler/test_handler_diffuser.py

diff --git a/tests/engine/handler/test_handler_diffuser.py b/tests/engine/handler/test_handler_diffuser.py
new file mode 100644
index 00000000..aba1fa11
--- /dev/null
+++ b/tests/engine/handler/test_handler_diffuser.py
@@ -0,0 +1,241 @@
+# Copyright 2025 - Pruna AI GmbH. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the DiffuserHandler inference handler."""
+
+from __future__ import annotations
+
+import inspect
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from datasets import Dataset
+from PIL import Image
+
+from pruna.config.smash_config import SmashConfig
+from pruna.data.collate import (
+    image_generation_collate,
+    prompt_collate,
+    prompt_with_auxiliaries_collate,
+)
+from pruna.data.pruna_datamodule import PrunaDataModule
+from pruna.engine.handler.handler_diffuser import DiffuserHandler
+from pruna.engine.pruna_model import PrunaModel
+from pruna.engine.utils import move_to_device
+
+
+def _img(size: int = 32) -> Image.Image:
+    """Return a random RGB PIL image."""
+    return Image.fromarray(np.random.randint(0, 255, (size, size, 3), dtype=np.uint8))
+
+
+def _fake_output(n: int, size: int = 16) -> SimpleNamespace:
+    """Return a fake diffusers output object exposing an ``images`` list of PIL images."""
+    return SimpleNamespace(images=[_img(size) for _ in range(n)])
+
+
+class PromptPipe:
+    """Dummy pipeline whose call signature exposes a ``prompt`` parameter."""
+
+    def __call__(self, prompt, negative_prompt=None, num_inference_steps=1, generator=None, **kwargs):  # noqa: D102
+        return _fake_output(len(prompt) if isinstance(prompt, list) else 1)
+
+
+class ArgsPipe:
+    """Dummy pipeline whose call signature exposes only variadic ``args``."""
+
+    def __call__(self, *args, generator=None):  # noqa: D102
+        return _fake_output(1)
+
+
+class UncondPipe:
+    """Dummy pipeline whose call signature exposes neither ``prompt`` nor ``args``."""
+
+    def __call__(self, num_inference_steps=1, generator=None):  # noqa: D102
+        return _fake_output(1)
+
+
+def _handler_for(pipe: Any) -> DiffuserHandler:
+    """Build a DiffuserHandler from a pipeline's call signature, mirroring production wiring."""
+    return DiffuserHandler(call_signature=inspect.signature(pipe.__call__))
+
+
+# -- prepare_inputs: prompt branch fed by real pruna collates --
+# Tests are available only for the collates appropriate for a diffuser pipeline.
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_prompt_collate():
+    """prompt_collate output is wrapped as a prompt dict for the model call."""
+    handler = _handler_for(PromptPipe())
+    data = [{"text": "a cat"}, {"text": "a dog"}]
+    batch = prompt_collate(data)
+    prepared = handler.prepare_inputs(batch)
+    assert prepared == {"prompt": ["a cat", "a dog"]}
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_prompt_with_auxiliaries_collate():
+    """prompt_with_auxiliaries_collate: prompts are wrapped and auxiliaries are dropped."""
+    handler = _handler_for(PromptPipe())
+    data = [
+        {"text": "a cat", "category": "animal"},
+        {"text": "a dog", "category": "animal"},
+    ]
+    batch = prompt_with_auxiliaries_collate(data)
+    prepared = handler.prepare_inputs(batch)
+    assert prepared == {"prompt": ["a cat", "a dog"]}
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_image_generation_collate():
+    """image_generation_collate texts become the prompt for the model call."""
+    handler = _handler_for(PromptPipe())
+    data = [{"image": _img(), "text": "a cat"}, {"image": _img(), "text": "a dog"}]
+    batch = image_generation_collate(data, img_size=32)
+    prepared = handler.prepare_inputs(batch)
+    assert prepared == {"prompt": ["a cat", "a dog"]}
+
+
+# -- prepare_inputs: custom negative-prompt path --
+
+def custom_negative_prompt_collate(data: Any) -> tuple[dict[str, list[str]], None]:
+    """Collate a dataset into a dict of prompt/negative_prompt lists as the model input."""
+    return {"prompt": [d["prompt"] for d in data], "negative_prompt": [d["negative_prompt"] for d in data]}, None
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_negative_prompt_via_datamodule():
+    """A custom collate emitting a dict input flows through PrunaDataModule and is preserved."""
+    handler = _handler_for(PromptPipe())
+    ds = Dataset.from_dict(
+        {"prompt": ["a cat", "a dog"], "negative_prompt": ["blurry", "low quality"]}
+    )
+    dm = PrunaDataModule(ds, ds, ds, custom_negative_prompt_collate, dataloader_args={})
+    batch = next(iter(dm.train_dataloader(batch_size=2, shuffle=False)))
+    prepared = handler.prepare_inputs(batch)
+    print(prepared)
+    assert prepared == {"prompt": ["a cat", "a dog"], "negative_prompt": ["blurry", "low quality"]}
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_dict_batch_preserved():
+    """A dict passed directly as the first batch element is forwarded unchanged."""
+    handler = _handler_for(PromptPipe())
+    payload = {"prompt": ["a cat"], "negative_prompt": ["blurry"]}
+    prepared = handler.prepare_inputs((payload, {}))
+    assert prepared == payload
+
+
+# ── prepare_inputs: other signature branches ─────────────────────
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_args_branch():
+    """When the signature only exposes *args, the first element is returned as-is."""
+    handler = _handler_for(ArgsPipe())
+    prepared = handler.prepare_inputs((["a cat", "a dog"], {}))
+    assert prepared == ["a cat", "a dog"]
+
+
+@pytest.mark.cpu
+def test_prepare_inputs_unconditional_branch():
+    """Unconditional pipelines (no prompt/args) receive no prepared inputs."""
+    handler = _handler_for(UncondPipe())
+    assert handler.prepare_inputs((["a cat"], {})) is None
+
+
+# ── process_output and __init__ ──────────────────────────────────
+
+
+@pytest.mark.cpu
+def test_process_output_stacks_images():
+    """process_output stacks PIL images into a uint8 (N, 3, H, W) tensor."""
+    handler = _handler_for(PromptPipe())
+    output = handler.process_output(_fake_output(2, size=16))
+    assert isinstance(output, torch.Tensor)
+    assert output.shape == (2, 3, 16, 16)
+    assert output.dtype == torch.uint8
+
+
+@pytest.mark.cpu
+def test_init_sets_generator_and_merges_model_args():
+    """__init__ provides a fixed-seed generator and merges extra model args without dropping it."""
+    handler = DiffuserHandler(
+        call_signature=inspect.signature(PromptPipe().__call__),
+        model_args={"num_inference_steps": 2},
+    )
+    assert isinstance(handler.model_args["generator"], torch.Generator)
+    assert handler.model_args["num_inference_steps"] == 2
+
+
+# ── real end-to-end generation on a tiny CPU model ───────────────
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("model_fixture", ["sd_tiny_random", "flux_tiny_random"], indirect=True)
+def test_generation_prompt_batch(model_fixture: tuple[Any, SmashConfig]):
+    """Real generation from a prompt-list batch returns a stacked image tensor."""
+    model, smash_config = model_fixture
+    smash_config.device = "cpu"
+    pruna_model = PrunaModel(model, smash_config=smash_config)
+    move_to_device(pruna_model, "cpu")
+    pruna_model.inference_handler.model_args["num_inference_steps"] = 2
+
+    batch = (["a red apple", "a blue car"], None)
+    outputs = pruna_model.run_inference(batch)
+
+    assert isinstance(outputs, torch.Tensor)
+    assert outputs.shape[0] == 2
+    assert outputs.shape[1] == 3
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("model_fixture", ["sd_tiny_random", "flux_tiny_random"], indirect=True)
+def test_generation_negative_prompt_batch(model_fixture: tuple[Any, SmashConfig]):
+    """Real generation with a negative_prompt dict batch reaches the pipeline and returns an image tensor."""
+    model, smash_config = model_fixture
+    smash_config.device = "cpu"
+    pruna_model = PrunaModel(model, smash_config=smash_config)
+    move_to_device(pruna_model, "cpu")
+    pruna_model.inference_handler.model_args["num_inference_steps"] = 2
+
+    batch = ({"prompt": ["a red apple"], "negative_prompt": ["blurry, low quality"]}, None)
+    outputs = pruna_model.run_inference(batch)
+
+    assert isinstance(outputs, torch.Tensor)
+    assert outputs.shape[0] == 1
+    assert outputs.shape[1] == 3
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("model_fixture", ["flux2_tiny_random"], indirect=True)
+def test_generation_flux2_image_prompt_batch(model_fixture: tuple[Any, SmashConfig]):
+    """Flux2 accepts an image+prompt dict batch end-to-end via the dict escape hatch."""
+    model, smash_config = model_fixture
+    smash_config.device = "cpu"
+    pruna_model = PrunaModel(model, smash_config=smash_config)
+    move_to_device(pruna_model, "cpu")
+    pruna_model.inference_handler.model_args["num_inference_steps"] = 2
+    pruna_model.inference_handler.model_args["text_encoder_out_layers"] = (1,)
+
+    batch = ({"prompt": ["a red apple"], "image": [_img(64)]}, None)
+    outputs = pruna_model.run_inference(batch)
+
+    assert isinstance(outputs, torch.Tensor)
+    assert outputs.shape[0] == 1
+    assert outputs.shape[1] == 3
diff --git a/tests/fixtures.py b/tests/fixtures.py
index d130b646..ae202717 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -206,6 +206,7 @@ def get_autoregressive_text_to_image_model(model_id: str) -> tuple[Any, SmashCon
     "flux_tiny_random_with_tokenizer": partial(
         get_diffusers_model_with_tokenizer, "katuni4ka/tiny-random-flux", torch_dtype=torch.float16
     ),
+    "flux2_tiny_random": partial(get_diffusers_model, "tiny-random/flux2", torch_dtype=torch.bfloat16),
     # text generation models
     "opt_tiny_random": partial(get_automodel_transformers, "yujiepan/opt-tiny-random"),
     "smollm_135m": partial(get_automodel_transformers, "HuggingFaceTB/SmolLM2-135M"),

From 413281546f1dc32ed1abb1f83fa4a7717da99969 Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Thu, 9 Jul 2026 08:58:37 +0000
Subject: [PATCH 5/6] style: minor formatting improvements in docs

---
 src/pruna/data/collate.py                    | 2 +-
 src/pruna/engine/handler/handler_diffuser.py | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/pruna/data/collate.py b/src/pruna/data/collate.py
index fadb1407..7a3bd955 100644
--- a/src/pruna/data/collate.py
+++ b/src/pruna/data/collate.py
@@ -255,7 +255,7 @@ def text_generation_collate(
     -------
     Tuple[torch.Tensor, torch.Tensor]
         A tuple with tokens from all text prompts as first element,
-        and their corresponding next tokens as second element. 
+        and their corresponding next tokens as second element.
     """
     text_col = _resolve_column(column_map, "text")
     input_ids = []
diff --git a/src/pruna/engine/handler/handler_diffuser.py b/src/pruna/engine/handler/handler_diffuser.py
index 33aaaf61..859430b4 100644
--- a/src/pruna/engine/handler/handler_diffuser.py
+++ b/src/pruna/engine/handler/handler_diffuser.py
@@ -59,7 +59,6 @@ def prepare_inputs(
         2. To provide additional arguments to the model (e.g., negative prompts), \
         construct the first element as a dictionary (e.g., ``{"prompt": [...], "negative_prompt": [...]}``).
 
-
         Parameters
         ----------
         batch : Tuple[List[str] | torch.Tensor | dict[str, Any], ...]

From 0cb2f2627c87db9ffa0cd33a61b35f3ba293e7e3 Mon Sep 17 00:00:00 2001
From: Kaloyan Fachikov <kaloyan.fachikov@pruna.ai>
Date: Thu, 9 Jul 2026 11:42:41 +0000
Subject: [PATCH 6/6] chore: update diffusers version in pyproject.toml

This change is necessary to ensure Flux2KleinPipeline is supported.
---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 96d660d1..152bf30d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -121,7 +121,7 @@ dependencies = [
     "datasets>=3.0",
     "numpy>=1.24.4",
     "numpydoc>=1.6.0",
-    "diffusers>=0.21.4",
+    "diffusers>=0.37.0",
     "torch_pruning",
     "ConfigSpace>=1.2.1",
     "sentencepiece",