Skip to content

Commit

Permalink
Add a way to set the timestep multiplier in the flow sampling.
Browse files Browse the repository at this point in the history
comfyanonymous committed Jul 6, 2024
1 parent ff63893 commit 2dc84d1
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions comfy/model_sampling.py
Original file line number Diff line number Diff line change
@@ -190,11 +190,12 @@ def __init__(self, model_config=None):
else:
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
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)

@property
@@ -206,10 +207,10 @@ def sigma_max(self):
return self.sigmas[-1]

def timestep(self, sigma):
return sigma * 1000
return sigma * self.multiplier

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):
if percent <= 0.0:

0 comments on commit 2dc84d1

Please sign in to comment.