Skip to content

Commit

Permalink
Only append zero to noise schedule if last sigma isn't zero.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Jul 20, 2024
1 parent 11b7414 commit 95fa954
Showing 1 changed file with 16 additions and 4 deletions.
20 changes: 16 additions & 4 deletions comfy/samplers.py
Original file line number Diff line number Diff line change
Expand Up @@ -313,30 +313,42 @@ def simple_scheduler(model_sampling, steps):
def ddim_scheduler(model_sampling, steps):
s = model_sampling
sigs = []
ss = max(len(s.sigmas) // steps, 1)
x = 1
if math.isclose(float(s.sigmas[x]), 0, abs_tol=0.00001):
steps += 1
sigs = []
else:
sigs = [0.0]

ss = max(len(s.sigmas) // steps, 1)
while x < len(s.sigmas):
sigs += [float(s.sigmas[x])]
x += ss
sigs = sigs[::-1]
sigs += [0.0]
return torch.FloatTensor(sigs)

def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
s = model_sampling
start = s.timestep(s.sigma_max)
end = s.timestep(s.sigma_min)

append_zero = True
if sgm:
timesteps = torch.linspace(start, end, steps + 1)[:-1]
else:
if math.isclose(float(s.sigma(end)), 0, abs_tol=0.00001):
steps += 1
append_zero = False
timesteps = torch.linspace(start, end, steps)

sigs = []
for x in range(len(timesteps)):
ts = timesteps[x]
sigs.append(s.sigma(ts))
sigs += [0.0]
sigs.append(float(s.sigma(ts)))

if append_zero:
sigs += [0.0]

return torch.FloatTensor(sigs)

# Implemented based on: https://arxiv.org/abs/2407.12173
Expand Down

0 comments on commit 95fa954

Please sign in to comment.