incorrect stochastic adjustment -- image goes to black when eta > 0

[dfl]

there are some differences in structure between Diffusers (more explicit time scheduling) and ComfyUI (array of sigmas, and abstracted scheduling), so I was unable to correctly translate. Perhaps someone with more knowledge about diffusers could assist?

[George0726]

Hi author,
I am trying to write the training code for TCD scheduler and I found the same problem images go all black. Have you ever solved the issue?

Many thanks

[dfl]

probably some values are going out of range, either blowing up or going to zero? I have not figured it out, but made a request at the TCD authors’ repository. Diffusers architecture is quite different than ComfyUI scheduling, I couldn’t figure out how to properly offset the timestep by eta.

[George0726]

Many thanks, I will check step by step.

[MoonRide303]

TCDScheduler node seems to be working (sample workflow that works pretty nicely below), but when trying to use SamplerTCD - always black output.

MoonRide SDXL-Lightning+TCD workflow v1.json

[laksjdjf]

What about this implementation?
You can get time from sigma by model.inner_model.inner_model.model_sampling.timestep(sigma)
and sigma from time by model.inner_model.inner_model.model_sampling.sigma(timestep)

@torch.no_grad()
def sampler_tcd(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, gamma=None):
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})

        d = to_d(x, sigmas[i], denoised)

        sigma_from = sigmas[i]
        sigma_to = sigmas[i + 1]

        t = model.inner_model.inner_model.model_sampling.timestep(sigma_from)
        down_t = (1 - gamma) * t
        sigma_down = model.inner_model.inner_model.model_sampling.sigma(down_t)

        if sigma_down > sigma_to:
            sigma_down = sigma_to

        sigma_up = (sigma_to ** 2 - sigma_down ** 2) ** 0.5
        
        # same as euler ancestral
        d = to_d(x, sigma_from, denoised)
        dt = sigma_down - sigma_from
        x = x + d * dt
        if sigma_to > 0:
            x = x + noise_sampler(sigma_from, sigma_to) * sigma_up
    return x

[dfl]

thanks @laksjdjf! This works, though somehow the quality is not as good as the reference implementation: https://huggingface.co/spaces/h1t/TCD

is it correct to get sigmas from ddim_uniform scheduler? I picked that because I see some DDIM references inside the reference diffusers.TCDScheduler code.

(edit: aha, I see now there should be some DDPM sampler behavior in the step, it is tricky to figure out how to normalize properly)

(As an aside, Intuitively I think it would make more sense to have gamma=0 be soft like euler, and gamma=1.0 for noisy details)

[josemerinom]

Hello, I did "git pull" and I have this error in comfyui

[dfl]

@josemerinom sorry I had a bad merge, please pull again