298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-03-15 14:09:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a way to set the timestep multiplier in the flow sampling.

Browse Source
This commit is contained in:
comfyanonymous 2024-07-06 04:06:03 -04:00
parent ff63893d10
commit 2dc84d1444
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/model_sampling.py
View File

@ -190,11 +190,12 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
else: else:
sampling_settings = {} sampling_settings = {}
self.set_parameters(shift=sampling_settings.get("shift", 1.0)) self.set_parameters(shift=sampling_settings.get("shift", 1.0), multiplier=sampling_settings.get("multiplier", 1000))
def set_parameters(self, shift=1.0, timesteps=1000): def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):
self.shift = shift self.shift = shift
ts = self.sigma(torch.arange(1, timesteps + 1, 1)) self.multiplier = multiplier
ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps) * multiplier)
self.register_buffer('sigmas', ts) self.register_buffer('sigmas', ts)
@property @property
@ -206,10 +207,10 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
return self.sigmas[-1] return self.sigmas[-1]
def timestep(self, sigma): def timestep(self, sigma):
return sigma * 1000 return sigma * self.multiplier
def sigma(self, timestep): def sigma(self, timestep):
return time_snr_shift(self.shift, timestep / 1000) return time_snr_shift(self.shift, timestep / self.multiplier)
def percent_to_sigma(self, percent): def percent_to_sigma(self, percent):
if percent <= 0.0: if percent <= 0.0: