feat: img-saver-extended

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-added-large-files
       - id: check-case-conflict
@@ -13,21 +13,45 @@ repos:
       - id: trailing-whitespace
         exclude: \.md$
       - id: no-commit-to-branch
+
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.12
+    rev: v0.14.6
     hooks:
       - id: ruff-check
         args: [--fix]
+        files: \.py$
       - id: ruff-format
+        files: \.py$
+
+  - repo: https://github.com/trufflesecurity/trufflehog
+    rev: v3.91.1
+    hooks:
+      - id: trufflehog
+        name: TruffleHog Secrets Scanner
+        entry: trufflehog
+        language: golang
+        types_or: [python, yaml, json, text]
+        args:
+          [
+            "filesystem",
+            "src",
+            "tests",
+            ".github/workflows",
+            "--results=verified,unknown",
+            "--exclude-paths=.venv",
+            "--fail"
+          ]
+        stages: ["pre-commit", "pre-push"]
 
   - repo: local
     hooks:
       - id: ty
-        name: ty check
+        name: type checking using ty
         entry: uvx ty check .
         language: system
         types: [python]
         pass_filenames: false
+        files: \.py$
 
   - repo: local
     hooks:
@@ -39,7 +63,7 @@ repos:
           grep -v "^D" |
           cut -f2- |
           while IFS= read -r file; do
-              if [ -f "$file" ] && ["$file" != ".pre-commit-config.yaml"] && grep -q "pruna_pro" "$file"; then
+              if [ -f "$file" ] && [ "$file" != ".pre-commit-config.yaml" ] && grep -q "pruna_pro" "$file"; then
                   echo "Error: pruna_pro found in staged file $file"
                   exit 1
               fi
@@ -48,10 +72,4 @@ repos:
         language: system
         stages: [pre-commit]
         types: [python]
-        exclude: "^docs/"
-      - id: trufflehog
-        name: TruffleHog
-        description: Detect secrets in your data.
-        entry: bash -c 'git diff --cached --name-only | xargs -I {} trufflehog filesystem {} --fail --no-update'
-        language: system
-        stages: ["pre-commit", "pre-push"]
+        files: \.py$

pyproject.toml

@@ -31,10 +31,13 @@ unsupported-operator = "ignore" # mypy supports | syntax with from __future__ im
 invalid-argument-type = "ignore" # mypy is more permissive with argument types
 invalid-return-type = "ignore" # mypy is more permissive with return types
 invalid-parameter-default = "ignore" # mypy is more permissive with parameter defaults
+possibly-missing-attribute = "ignore" # mypy is more permissive with attribute access
+possibly-unbound-attribute = "ignore"
+possibly-missing-import = "ignore" # mypy is more permissive with imports
 no-matching-overload = "ignore" # mypy is more permissive with overloads
 unresolved-reference = "ignore" # mypy is more permissive with references
-possibly-unbound-import = "ignore"
 missing-argument = "ignore"
+possibly-unbound-import = "ignore"
 
 [tool.coverage.run]
 source = ["src/pruna"]
@@ -75,6 +78,7 @@ gptqmodel = [
   { index = "pruna_internal", marker = "sys_platform != 'darwin' or platform_machine != 'arm64'"},
   { index = "pypi", marker = "sys_platform == 'darwin' and platform_machine == 'arm64'"},
 ]
+clip = {git = "https://github.com/openai/CLIP.git", rev = "dcba3cb2e2827b402d2701e7e1c7d9fed8a20ef1"}
 
 [project]
 name = "pruna"
@@ -186,6 +190,7 @@ dev = [
 ]
 cpu = []
 
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/pruna/data/utils.py

@@ -184,7 +184,7 @@ def recover_text_from_dataloader(dataloader: DataLoader, tokenizer: Any) -> list
     return texts
 
 
-def stratify_dataset(dataset: Dataset, sample_size: int, seed: int = 42) -> Dataset:
+def stratify_dataset(dataset: Dataset, sample_size: int, seed: int = 42, partition_strategy: str = "random", partition_index: int = 0) -> Dataset:
     """
     Stratify the dataset into a specific size.
 
@@ -196,6 +196,10 @@ def stratify_dataset(dataset: Dataset, sample_size: int, seed: int = 42) -> Data
         The size to stratify.
     seed : int
         The seed to use for sampling the dataset.
+    partition_strategy : str
+        The strategy to use for partitioning the dataset. Can be "indexed" or "random".
+    partition_index : int
+        The index to use for partitioning the dataset.
 
     Returns
     -------
@@ -211,8 +215,13 @@ def stratify_dataset(dataset: Dataset, sample_size: int, seed: int = 42) -> Data
         return dataset
 
     indices = list(range(dataset_length))
-    random.Random(seed).shuffle(indices)
-    selected_indices = indices[:sample_size]
+    if partition_strategy == "indexed":
+        selected_indices = indices[sample_size*partition_index:sample_size*(partition_index+1)]
+    elif partition_strategy == "random":
+        random.Random(seed).shuffle(indices)
+        selected_indices = indices[:sample_size]
+    else:
+        raise ValueError(f"Invalid partition strategy: {partition_strategy}")
     dataset = dataset.select(selected_indices)
     return dataset
 

src/pruna/engine/handler/handler_diffuser.py

@@ -15,10 +15,9 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import torch
-from torchvision import transforms
 
 from pruna.engine.handler.handler_inference import InferenceHandler
 from pruna.logging.logger import pruna_logger
@@ -28,10 +27,6 @@ class DiffuserHandler(InferenceHandler):
     """
     Handle inference arguments, inputs and outputs for diffusers models.
 
-    A generator with a fixed seed (42) is passed as an argument to the model for reproducibility.
-    The first element of the batch is passed as input to the model.
-    The generated outputs are expected to have .images attribute.
-
     Parameters
     ----------
     call_signature : inspect.Signature
@@ -40,12 +35,18 @@ class DiffuserHandler(InferenceHandler):
         The arguments to pass to the model.
     """
 
-    def __init__(self, call_signature: inspect.Signature, model_args: Optional[Dict[str, Any]] = None) -> None:
-        default_args = {"generator": torch.Generator("cpu").manual_seed(42)}
+    def __init__(
+        self,
+        call_signature: inspect.Signature,
+        model_args: Optional[Dict[str, Any]] = None,
+        seed_strategy: Literal["per_sample", "no_seed"] = "no_seed",
+        global_seed: int | None = None,
+    ) -> None:
         self.call_signature = call_signature
-        if model_args:
-            default_args.update(model_args)
-        self.model_args = default_args
+        self.model_args = model_args if model_args else {}
+        # We want the default output type to be pytorch tensors.
+        self.model_args["output_type"] = "pt"
+        self.configure_seed(seed_strategy, global_seed)
 
     def prepare_inputs(
         self, batch: List[str] | torch.Tensor | Tuple[List[str] | torch.Tensor | dict[str, Any], ...] | dict[str, Any]
@@ -83,13 +84,36 @@ def process_output(self, output: Any) -> torch.Tensor:
         torch.Tensor
             The processed images.
         """
-        generated = output.images
-        return torch.stack([transforms.PILToTensor()(g) for g in generated])
+        if hasattr(output, "images"):
+            generated = output.images
+        #  For video models.
+        elif hasattr(output, "frames"):
+            generated = output.frames
+        else:
+            # Maybe the user is calling the pipeline with return_dict = False,
+            # which then returns the generated image / video in a tuple
+            generated = output[0]
+        return generated.float()
 
     def log_model_info(self) -> None:
         """Log information about the inference handler."""
         pruna_logger.info(
-            "Detected diffusers model. Using DiffuserHandler with fixed seed.\n"
-            "- The first element of the batch is passed as input.\n"
-            "- The generated outputs are expected to have .images attribute."
+            "Detected diffusers model. Using DiffuserHandler.\n- The first element of the batch is passed as input.\n"
+            "Inference outputs are expected to have either have an `images` attribute or a `frames` attribute."
+            "Or be a tuple with the generated image / video as the first element."
         )
+
+    def set_seed(self, seed: int) -> None:
+        """
+        Set the random seed for the current process.
+
+        Parameters
+        ----------
+        seed : int
+            The seed to set.
+        """
+        self.model_args["generator"] = torch.Generator("cpu").manual_seed(seed)
+
+    def remove_seed(self) -> None:
+        """Remove the seed from the current process."""
+        self.model_args["generator"] = None

src/pruna/engine/handler/handler_inference.py

@@ -14,9 +14,11 @@
 
 from __future__ import annotations
 
+import random
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Literal, Tuple
 
+import numpy as np
 import torch
 
 from pruna.data.utils import move_batch_to_device
@@ -98,3 +100,75 @@ def move_inputs_to_device(
             return move_batch_to_device(inputs, device)
         except torch.cuda.OutOfMemoryError as e:
             raise e
+
+    def configure_seed(self, seed_strategy: Literal["per_sample", "no_seed"], global_seed: int | None) -> None:
+        """
+        Set the random seed according to the chosen strategy.
+
+        - If `seed_strategy="per_sample"`,the `global_seed` is used as a base to derive a different seed for each
+        sample. This ensures reproducibility while still producing variation across samples,
+        making it the preferred option for benchmarking.
+        - If `seed_strategy="no_seed"`, no seed is set internally.
+        The user is responsible for managing seeds if reproducibility is required.
+
+        Parameters
+        ----------
+        seed_strategy : Literal["per_sample", "no_seed"]
+            The seeding strategy to apply.
+        global_seed : int | None
+            The base seed value to use (if applicable).
+        """
+        self.seed_strategy = seed_strategy
+        validate_seed_strategy(seed_strategy, global_seed)
+        if global_seed is not None:
+            self.global_seed = global_seed
+            self.set_seed(global_seed)
+        else:
+            self.remove_seed()
+
+    def set_seed(self, seed: int) -> None:
+        """
+        Set the random seed for the current process.
+
+        Parameters
+        ----------
+        seed : int
+            The seed to set.
+        """
+        #  With the default handler, we can't assume anything about the model,
+        #  so we are setting the seed for all RNGs available.
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+
+    def remove_seed(self) -> None:
+        """Remove the seed from the current process."""
+        random.seed(None)
+        np.random.seed(None)
+        #  We can't really remove the seed from the PyTorch RNG, so we are reseeding with torch.seed().
+        #  torch.seed() creates a non-deterministic random number.
+        torch.manual_seed(torch.seed())
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(torch.seed())
+
+
+def validate_seed_strategy(seed_strategy: Literal["per_sample", "no_seed"], global_seed: int | None) -> None:
+    """
+    Check the consistency of the seed strategy and the global seed.
+
+    If the seed strategy is "no_seed", the global seed must be None.
+    If the seed strategy is or "per_sample", the user must provide a global seed.
+
+    Parameters
+    ----------
+    seed_strategy : Literal["per_sample", "no_seed"]
+        The seeding strategy to apply.
+    global_seed : int | None
+        The base seed value to use (if applicable).
+    """
+    if seed_strategy != "no_seed" and global_seed is None:
+        raise ValueError("Global seed must be provided if seed strategy is not 'no_seed'.")
+    elif global_seed is not None and seed_strategy == "no_seed":
+        raise ValueError("Seed strategy cannot be 'no_seed' if global seed is provided.")

src/pruna/engine/pruna_model.py

@@ -24,7 +24,7 @@
 
 from pruna.config.smash_config import SmashConfig
 from pruna.engine.handler.handler_utils import register_inference_handler
-from pruna.engine.load import load_pruna_model, load_pruna_model_from_pretrained
+from pruna.engine.load import filter_load_kwargs, load_pruna_model, load_pruna_model_from_pretrained
 from pruna.engine.save import save_pruna_model, save_pruna_model_to_hub
 from pruna.engine.utils import get_device, get_nn_modules, set_to_eval
 from pruna.logging.filter import apply_warning_filter
@@ -108,6 +108,8 @@ def run_inference(self, batch: Any) -> Any:
             )
         inference_function = getattr(self, inference_function_name)
 
+        self.inference_handler.model_args = filter_load_kwargs(self.model.__call__, self.inference_handler.model_args)
+
         if prepared_inputs is None:
             outputs = inference_function(**self.inference_handler.model_args)
         elif isinstance(prepared_inputs, dict):