refactor: configure seeding and tests

Author: begumcig

pyproject.toml

@@ -35,6 +35,7 @@ possibly-missing-attribute = "ignore" # mypy is more permissive with attribute a
 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
+missing-argument = "ignore"
 
 [tool.coverage.run]
 source = ["src/pruna"]

src/pruna/engine/handler/handler_diffuser.py

@@ -19,7 +19,7 @@
 
 import torch
 
-from pruna.engine.handler.handler_inference import InferenceHandler, validate_seed_strategy
+from pruna.engine.handler.handler_inference import InferenceHandler
 from pruna.logging.logger import pruna_logger
 
 
@@ -93,35 +93,27 @@ def process_output(self, output: Any) -> torch.Tensor:
             # 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
+        return generated.float()
 
     def log_model_info(self) -> None:
         """Log information about the inference handler."""
         pruna_logger.info(
             "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 configure_seed(self, seed_strategy: Literal["per_sample", "no_seed"], global_seed: int | None) -> None:
+    def set_seed(self, seed: int) -> 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.
+        Set the random seed for the current process.
 
         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).
+        seed : int
+            The seed to set.
         """
-        self.seed_strategy = seed_strategy
-        validate_seed_strategy(seed_strategy, global_seed)
-        if global_seed is not None:
-            self.global_seed = global_seed
-            self.model_args["generator"] = torch.Generator("cpu").manual_seed(global_seed)
-        else:
-            self.model_args["generator"] = None  # Remove the seed.
+        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

@@ -122,38 +122,36 @@ def configure_seed(self, seed_strategy: Literal["per_sample", "no_seed"], global
         validate_seed_strategy(seed_strategy, global_seed)
         if global_seed is not None:
             self.global_seed = global_seed
-            set_seed(global_seed)
+            self.set_seed(global_seed)
         else:
-            remove_seed()
+            self.remove_seed()
 
+    def set_seed(self, seed: int) -> None:
+        """
+        Set the random seed for the current process.
 
-def set_seed(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() -> 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())
+        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:

tests/engine/test_handler.py

@@ -1,11 +1,11 @@
-import types
+import numpy as np
 import pytest
 import torch
 from pruna.engine.handler.handler_inference import (
-    set_seed,
     validate_seed_strategy,
 )
 from pruna.engine.handler.handler_diffuser import DiffuserHandler
+from pruna.engine.handler.handler_standard import StandardHandler
 from pruna.engine.pruna_model import PrunaModel
 from pruna.engine.utils import move_to_device
 #  Default handler tests, mainly for checking seeding.
@@ -25,14 +25,16 @@ def test_validate_seed_strategy_invalid(strategy, seed):
     with pytest.raises(ValueError):
         validate_seed_strategy(strategy, seed)
 
-
 def test_set_seed_reproducibility():
-    '''Test to see set_seed is reproducible'''
-    set_seed(42)
-    a = torch.randn(3)
-    set_seed(42)
-    b = torch.randn(3)
-    assert torch.equal(a, b)
+    inference_handler = StandardHandler()
+    inference_handler.set_seed(42)
+    torch_random_tensor = torch.randn(3)
+    numpy_random_tensor = np.random.randn(3)
+    inference_handler.set_seed(42)
+    torch_expected = torch.randn(3)
+    numpy_expected = np.random.randn(3)
+    assert torch.equal(torch_random_tensor, torch_expected)
+    assert np.array_equal(numpy_random_tensor, numpy_expected)
 
 
 # Diffuser handler tests, checking output processing and seeding.