From 3f8df52701a6d887089bef1e6f47a19327fdd680 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Thu, 25 Jan 2024 10:00:36 +0000 Subject: [PATCH] Correct sigmas cpu settings --- src/diffusers/schedulers/scheduling_consistency_models.py | 4 ++-- src/diffusers/schedulers/scheduling_deis_multistep.py | 4 ++-- src/diffusers/schedulers/scheduling_dpmsolver_multistep.py | 4 ++-- .../schedulers/scheduling_dpmsolver_multistep_inverse.py | 4 ++-- src/diffusers/schedulers/scheduling_dpmsolver_sde.py | 4 ++-- src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py | 4 ++-- .../schedulers/scheduling_euler_ancestral_discrete.py | 4 ++-- src/diffusers/schedulers/scheduling_euler_discrete.py | 4 ++-- src/diffusers/schedulers/scheduling_heun_discrete.py | 4 ++-- .../schedulers/scheduling_k_dpm_2_ancestral_discrete.py | 4 ++-- src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py | 4 ++-- src/diffusers/schedulers/scheduling_lms_discrete.py | 4 ++-- src/diffusers/schedulers/scheduling_sasolver.py | 4 ++-- src/diffusers/schedulers/scheduling_unipc_multistep.py | 4 ++-- 14 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_consistency_models.py b/src/diffusers/schedulers/scheduling_consistency_models.py index b9a21f9bbd37..e7e0dcbdc31e 100644 --- a/src/diffusers/schedulers/scheduling_consistency_models.py +++ b/src/diffusers/schedulers/scheduling_consistency_models.py @@ -98,7 +98,7 @@ def __init__( self.custom_timesteps = False self.is_scale_input_called = False self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication def index_for_timestep(self, timestep, schedule_timesteps=None): if schedule_timesteps is None: @@ -231,7 +231,7 @@ def set_timesteps( self.timesteps = torch.from_numpy(timesteps).to(device=device) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Modified _convert_to_karras implementation that takes in ramp as argument def _convert_to_karras(self, ramp): diff --git a/src/diffusers/schedulers/scheduling_deis_multistep.py b/src/diffusers/schedulers/scheduling_deis_multistep.py index 572078a9d604..e8bd5f8f68d4 100644 --- a/src/diffusers/schedulers/scheduling_deis_multistep.py +++ b/src/diffusers/schedulers/scheduling_deis_multistep.py @@ -187,7 +187,7 @@ def __init__( self.model_outputs = [None] * solver_order self.lower_order_nums = 0 self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def step_index(self): @@ -255,7 +255,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 8b60c02929d6..d70d4eec9b3e 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -227,7 +227,7 @@ def __init__( self.model_outputs = [None] * solver_order self.lower_order_nums = 0 self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def step_index(self): @@ -311,7 +311,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py index aea3c4c4b9f6..03fc3677d07f 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py @@ -213,7 +213,7 @@ def __init__( self.model_outputs = [None] * solver_order self.lower_order_nums = 0 self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication self.use_karras_sigmas = use_karras_sigmas @property @@ -294,7 +294,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py index a999a8adbfa7..20c294f95bd6 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py @@ -198,7 +198,7 @@ def __init__( self.noise_sampler = None self.noise_sampler_seed = noise_sampler_seed self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler.index_for_timestep def index_for_timestep(self, timestep, schedule_timesteps=None): @@ -348,7 +348,7 @@ def set_timesteps( self.mid_point_sigma = None self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication self.noise_sampler = None # for exp beta schedules, such as the one for `pipeline_shap_e.py` diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py index f4ee8b8bdc5d..f664374a4238 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py @@ -210,7 +210,7 @@ def __init__( self.sample = None self.order_list = self.get_order_list(num_train_timesteps) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication def get_order_list(self, num_inference_steps: int) -> List[int]: """ @@ -315,7 +315,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index ca188378a38f..acad67847237 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -216,7 +216,7 @@ def __init__( self.is_scale_input_called = False self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def init_noise_sigma(self): @@ -300,7 +300,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.timesteps = torch.from_numpy(timesteps).to(device=device) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index def _init_step_index(self, timestep): diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index c72f7ff336aa..6ed28f410aea 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -237,7 +237,7 @@ def __init__( self.use_karras_sigmas = use_karras_sigmas self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def init_noise_sigma(self): @@ -342,7 +342,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication def _sigma_to_t(self, sigma, log_sigmas): # get log sigma diff --git a/src/diffusers/schedulers/scheduling_heun_discrete.py b/src/diffusers/schedulers/scheduling_heun_discrete.py index d06459e0a264..a1ea18dcf168 100644 --- a/src/diffusers/schedulers/scheduling_heun_discrete.py +++ b/src/diffusers/schedulers/scheduling_heun_discrete.py @@ -148,7 +148,7 @@ def __init__( self.use_karras_sigmas = use_karras_sigmas self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication def index_for_timestep(self, timestep, schedule_timesteps=None): if schedule_timesteps is None: @@ -270,7 +270,7 @@ def set_timesteps( self.dt = None self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # (YiYi Notes: keep this for now since we are keeping add_noise function which use index_for_timestep) # for exp beta schedules, such as the one for `pipeline_shap_e.py` diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py index 523b1f4f3b96..4a1cdb561cea 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py @@ -140,7 +140,7 @@ def __init__( # set all values self.set_timesteps(num_train_timesteps, None, num_train_timesteps) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler.index_for_timestep def index_for_timestep(self, timestep, schedule_timesteps=None): @@ -300,7 +300,7 @@ def set_timesteps( self._index_counter = defaultdict(int) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t def _sigma_to_t(self, sigma, log_sigmas): diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py index e1e5124d70e5..57062c0d3586 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py @@ -140,7 +140,7 @@ def __init__( self.set_timesteps(num_train_timesteps, None, num_train_timesteps) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler.index_for_timestep def index_for_timestep(self, timestep, schedule_timesteps=None): @@ -285,7 +285,7 @@ def set_timesteps( self._index_counter = defaultdict(int) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def state_in_first_order(self): diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index a78fa0e42639..f5f52b06bd43 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -168,7 +168,7 @@ def __init__( self.is_scale_input_called = False self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def init_noise_sigma(self): @@ -280,7 +280,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to(device=device) self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication self.derivatives = [] diff --git a/src/diffusers/schedulers/scheduling_sasolver.py b/src/diffusers/schedulers/scheduling_sasolver.py index e25178fe8eb2..53c0dd2c8e16 100644 --- a/src/diffusers/schedulers/scheduling_sasolver.py +++ b/src/diffusers/schedulers/scheduling_sasolver.py @@ -212,7 +212,7 @@ def __init__( self.lower_order_nums = 0 self.last_sample = None self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def step_index(self): @@ -283,7 +283,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index c147e0142a32..1223213c69f3 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -198,7 +198,7 @@ def __init__( self.solver_p = solver_p self.last_sample = None self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property def step_index(self): @@ -269,7 +269,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic # add an index counter for schedulers that allow duplicated timesteps self._step_index = None - self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: