fix: parameters order assumption in diffuser

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

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

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 = {}