[Refactor] Remove set_seed (huggingface#289)

* [Refactor] Remove set_seed and class attributes

* apply anton's suggestiosn

* fix

* Apply suggestions from code review

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

* up

* update

* make style

* Apply suggestions from code review

Co-authored-by: Anton Lozhkov <anton@huggingface.co>

* make fix-copies

* make style

* make style and new copies

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
Co-authored-by: Anton Lozhkov <anton@huggingface.co>

Author: 3 people

pipelines/ddpm/pipeline_ddpm.py

@@ -56,7 +56,7 @@ def __call__(self, batch_size=1, generator=None, output_type="pil", **kwargs):
             model_output = self.unet(image, t)["sample"]
 
             # 2. compute previous image: x_t -> t_t-1
-            image = self.scheduler.step(model_output, t, image)["prev_sample"]
+            image = self.scheduler.step(model_output, t, image, generator=generator)["prev_sample"]
 
         image = (image / 2 + 0.5).clamp(0, 1)
         image = image.cpu().permute(0, 2, 3, 1).numpy()

pipelines/score_sde_ve/pipeline_score_sde_ve.py

@@ -30,7 +30,7 @@ def __call__(self, batch_size=1, num_inference_steps=2000, generator=None, outpu
 
         model = self.unet
 
-        sample = torch.randn(*shape) * self.scheduler.config.sigma_max
+        sample = torch.randn(*shape, generator=generator) * self.scheduler.config.sigma_max
         sample = sample.to(self.device)
 
         self.scheduler.set_timesteps(num_inference_steps)
@@ -42,11 +42,11 @@ def __call__(self, batch_size=1, num_inference_steps=2000, generator=None, outpu
             # correction step
             for _ in range(self.scheduler.correct_steps):
                 model_output = self.unet(sample, sigma_t)["sample"]
-                sample = self.scheduler.step_correct(model_output, sample)["prev_sample"]
+                sample = self.scheduler.step_correct(model_output, sample, generator=generator)["prev_sample"]
 
             # prediction step
             model_output = model(sample, sigma_t)["sample"]
-            output = self.scheduler.step_pred(model_output, t, sample)
+            output = self.scheduler.step_pred(model_output, t, sample, generator=generator)
 
             sample, sample_mean = output["prev_sample"], output["prev_sample_mean"]
 

pipelines/stochatic_karras_ve/pipeline_stochastic_karras_ve.py

@@ -19,6 +19,7 @@ class KarrasVePipeline(DiffusionPipeline):
     differential equations." https://arxiv.org/abs/2011.13456
     """
 
+    # add type hints for linting
     unet: UNet2DModel
     scheduler: KarrasVeScheduler
 

schedulers/scheduling_sde_ve.py

@@ -14,8 +14,8 @@
 
 # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
 
-# TODO(Patrick, Anton, Suraj) - make scheduler framework indepedent and clean-up a bit
-from typing import Union
+import warnings
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -98,6 +98,11 @@ def get_adjacent_sigma(self, timesteps, t):
         raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
 
     def set_seed(self, seed):
+        warnings.warn(
+            "The method `set_seed` is deprecated and will be removed in version `0.4.0`. Please consider passing a"
+            " generator instead.",
+            DeprecationWarning,
+        )
         tensor_format = getattr(self, "tensor_format", "pt")
         if tensor_format == "np":
             np.random.seed(seed)
@@ -111,14 +116,14 @@ def step_pred(
         model_output: Union[torch.FloatTensor, np.ndarray],
         timestep: int,
         sample: Union[torch.FloatTensor, np.ndarray],
-        seed=None,
+        generator: Optional[torch.Generator] = None,
+        **kwargs,
     ):
         """
         Predict the sample at the previous timestep by reversing the SDE.
         """
-        if seed is not None:
-            self.set_seed(seed)
-        # TODO(Patrick) non-PyTorch
+        if "seed" in kwargs and kwargs["seed"] is not None:
+            self.set_seed(kwargs["seed"])
 
         if self.timesteps is None:
             raise ValueError(
@@ -140,7 +145,7 @@ def step_pred(
         drift = drift - diffusion[:, None, None, None] ** 2 * model_output
 
         #  equation 6: sample noise for the diffusion term of
-        noise = self.randn_like(sample)
+        noise = self.randn_like(sample, generator=generator)
         prev_sample_mean = sample - drift  # subtract because `dt` is a small negative timestep
         # TODO is the variable diffusion the correct scaling term for the noise?
         prev_sample = prev_sample_mean + diffusion[:, None, None, None] * noise  # add impact of diffusion field g
@@ -151,14 +156,15 @@ def step_correct(
         self,
         model_output: Union[torch.FloatTensor, np.ndarray],
         sample: Union[torch.FloatTensor, np.ndarray],
-        seed=None,
+        generator: Optional[torch.Generator] = None,
+        **kwargs,
     ):
         """
         Correct the predicted sample based on the output model_output of the network. This is often run repeatedly
         after making the prediction for the previous timestep.
         """
-        if seed is not None:
-            self.set_seed(seed)
+        if "seed" in kwargs and kwargs["seed"] is not None:
+            self.set_seed(kwargs["seed"])
 
         if self.timesteps is None:
             raise ValueError(
@@ -167,7 +173,7 @@ def step_correct(
 
         # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
         # sample noise for correction
-        noise = self.randn_like(sample)
+        noise = self.randn_like(sample, generator=generator)
 
         # compute step size from the model_output, the noise, and the snr
         grad_norm = self.norm(model_output)

utils/dummy_scipy_objects.py

@@ -1,5 +1,6 @@
 # This file is autogenerated by the command `make fix-copies`, do not edit.
 # flake8: noqa
+
 from ..utils import DummyObject, requires_backends
 
 

utils/dummy_transformers_objects.py

@@ -1,5 +1,6 @@
 # This file is autogenerated by the command `make fix-copies`, do not edit.
 # flake8: noqa
+
 from ..utils import DummyObject, requires_backends