[Feature Request]: Trajectory consistency distillation sampler

[George0726]

Is there an existing issue for this?

• [X] I have searched the existing issues and checked the recent builds/commits

What would your feature do ?

Currently Trajectory consistency distillation sampler has achieved great performance on fast inference of diffusion models. It would be great if A1111 could support the TCD sampler for further usage. https://github.com/jabir-zheng/TCD

Proposed workflow

1. Go to Sampler.
2. Press TCD sampler.
3. change the related CFG score and steps, i.e. CFG=1.0; steps=4 for inference.

Additional information

No response

[l0stl0rd]

would love to see support for this.

[AviceProjects]

I would love to see support for this too.

[bigmover]

SO COOL！ The feature will be widely welcomed and enjoyed by the PEOPLE！ Any progress or need help? Reference https://huggingface.co/docs/diffusers/using-diffusers/inference_with_tcd_lora

[bigmover]

Is there an existing issue for this?

• [x] I have searched the existing issues and checked the recent builds/commits

What would your feature do ?

Currently Trajectory consistency distillation sampler has achieved great performance on fast inference of diffusion models. It would be great if A1111 could support the TCD sampler for further usage. https://github.com/jabir-zheng/TCD

Proposed workflow

1. Go to Sampler.
2. Press TCD sampler.
3. change the related CFG score and steps, i.e. CFG=1.0; steps=4 for inference.

Additional information

No response

Hi guys! @George0726 Any plan to support TCD? Do you need a teamer?

question: Why does Euler's method combined with TCD achieve better results than the TCD sampler?

[Roveer]

would love to see support for TCD !!!!

[George0726]

I have written the TCD EULER A sampler for the v1.6.0 for your reference. The sampler cannot work with V.1.10

def sample_tcd_euler_a(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, gamma=0.3):
    print(f"################################## USE gamma: {gamma}")

    def append_dims(x, target_dims):
        """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
        dims_to_append = target_dims - x.ndim
        if dims_to_append < 0:
            raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
        expanded = x[(...,) + (None,) * dims_to_append]
        # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
        # https://github.com/pytorch/pytorch/issues/84364
        return expanded.detach().clone() if expanded.device.type == 'mps' else expanded

    def to_d(x, sigma, denoised):
        """Converts a denoiser output to a Karras ODE derivative."""
        return (x - denoised) / append_dims(sigma, x.ndim)
    def sigma_to_t( sigma):
        return sigma.atan() / math.pi * 2
    def t_to_sigma(t):
        return (t * math.pi / 2).tan()
    def noise_scaling(sigma, noise, latent_image, max_denoise=False):
        if max_denoise:
            noise = noise * torch.sqrt(1.0 + sigma ** 2.0)
        else:
            noise = noise * sigma

        noise += latent_image
        return noise

    extra_args = {} if extra_args is None else extra_args
    noise_sampler = k_diffusion.sampling.default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])
    for i in tqdm.auto.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)
        t = sigma_to_t(sigma_from)
        down_t = (1 - gamma) * t
        # sigma_down = model.inner_model.inner_model.model_sampling.sigma(down_t)
        sigma_down = t_to_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 += d * dt

        if sigma_to > 0 and gamma > 0:
            x = noise_scaling(sigma_up, noise_sampler(sigma_from, sigma_to), x)
    return x