feat: add tcd sampler

[JettHu]

This PR solved #2985 . Add Trajectory Consistency Distillation Support.

• Implemented a new sampling function sample_tcd.
• TCD use alphas_cumprod to calculate_denoised & noise_scaling, so I add register_buffer alphas_cumprod in ModelSamplingDiscrete.
• Add a tcd sampling in comfy_extras.nodes_model_advanced.ModelSamplingDiscrete
• Add a SamplerTCD node in comfy_extras.nodes_custom_sampler for set eta (referred to as gamma in the paper).

The upper part of the screenshot below shows setting eta (gamma in paper) value through SamplerTCD. In the lower part, selecting tcd directly in KSampler will use the default 0.3 as gamma.

TCD, significantly improves image quality at low NFEs (Number of Function Evaluations, or steps), and can produce more detailed results than the teacher model at high NFEs. Just like lcm, it can be directly applied to various models, and combined with LoRA . Although tcd-lora can directly use common PF-ODE solvers (such as DDIM or SDE-DPM-Solver), implementing tcd-scheduler can bring better quality.

Diffusers library already integrated TCDScheduler.

If someone are in a hurry, you can use my plug-in repository.

[MoonRide303]

@JettHu it seems to be just a sampler in this PR, I don't see TCD scheduler included - am I missing something? Isn't TCD scheduler necessary too, for TCD method to work as intended?

I've tried TCD scheduler available as custom node from here: https://github.com/dfl/comfyui-tcd-scheduler, and I've got different results from sgm_uniform.

Just this PR:

This PR + TCD scheduler (custom node):

Test workflow:
SDXL + TCD workflow.json

UPDATE:
For lower number of steps (and eta) sgm_uniform seems to be working better than TCD sampler:

vs

And for higher number of steps (and eta) results seem to be getting similar:

vs

So... it seems TCD sampler with sgm_uniform scheduler should do.

@comfyanonymous?

[rabidcopy]

!!! Exception during processing !!!
Traceback (most recent call last):
  File "C:\Users\user\Desktop\comfyui\execution.py", line 151, in recursive_execute
    output_data, output_ui = get_output_data(obj, input_data_all)
  File "C:\Users\user\Desktop\comfyui\execution.py", line 81, in get_output_data
    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
  File "C:\Users\user\Desktop\comfyui\execution.py", line 74, in map_node_over_list
    results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
  File "C:\Users\user\Desktop\comfyui\comfy_extras\nodes_custom_sampler.py", line 392, in sample
    samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
  File "C:\Users\user\Desktop\comfyui\comfy\sample.py", line 42, in sample_custom
    samples = comfy.samplers.sample(model, noise, positive, negative, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
  File "C:\Users\user\Desktop\comfyui\comfy\samplers.py", line 657, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "C:\Users\user\Desktop\comfyui\comfy\samplers.py", line 644, in sample
    output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "C:\Users\user\Desktop\comfyui\comfy\samplers.py", line 623, in inner_sample
    samples = sampler.sample(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
  File "C:\Users\user\Desktop\comfyui\comfy\samplers.py", line 534, in sample
    samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
  File "C:\Users\user\Desktop\comfyui\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "C:\Users\user\Desktop\comfyui\comfy\k_diffusion\sampling.py", line 830, in sample_tcd
    alpha_prod_s = model_sampling.alphas_cumprod[timesteps_s]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

This came up when trying to us this PR while the original repo works fine, oddly. (Am several commits behind though..)

[JettHu]

@JettHu it seems to be just a sampler in this PR, I don't see TCD scheduler included - am I missing something?

Sorry, it was my mistake. this PR does contain only a sampler. Because it's called scheduler in diffusers, I'm used to it. So the title of this PR should be called tcd sampler, which is more reasonable.

Isn't TCD scheduler necessary too, for TCD method to work as intended?
I've tried TCD scheduler available as custom node from here: https://github.com/dfl/comfyui-tcd-scheduler, and I've got different results from sgm_uniform.

TCD and LCM use the same scheduler, which are simple and sgm_uniform mentioned here. No additional TCD scheduler is required. After my testing, the sigmas of scheduler simple is consistent with diffusers. However both can give good results.

I have not seen the implementation details of comfyui-tcd-scheduler. As mentioned here, my implementation ComfyUI-TCD may be more consistent with the original TCD. The content of this PR basically comes from my repo comfyui-tcd ComfyUI-TCD.

[JettHu]

@rabidcopy Do you mind sharing your workflow and comfyui startup parameters?

I updated a version to ensure that the indices in sample_tcd are on the cpu.

[rabidcopy]

@rabidcopy Do you mind sharing your workflow and comfyui startup parameters?

I updated a version to ensure that the indices in sample_tcd are on the cpu.

Thanks for the update. The PR seems to work now without throwing any errors. I was using the workflow from here changed to use the node from this PR instead of your dedicated node for it. I think the device issue might have stemmed from my current usage of the DirectML backend as I am currently on Windows with an AMD GPU. I think the reason the dedicated node didn't have this issue for me is due to this line right here. https://github.com/JettHu/ComfyUI-TCD/blob/main/__init__.py#L61 But that's just a guess. Edit: Also for further context, I was inferencing an SDXL model that was loaded in low VRAM mode.

[RandomGitUser321]

@JettHu it seems to be just a sampler in this PR, I don't see TCD scheduler included - am I missing something? Isn't TCD scheduler necessary too, for TCD method to work as intended?

I've tried TCD scheduler available as custom node from here: https://github.com/dfl/comfyui-tcd-scheduler, and I've got different results from sgm_uniform.

So... it seems TCD sampler with sgm_uniform scheduler should do.

From the link you provided:
"NOTE: SamplerTCD is a WIP and currently just operates as DDIM with no gamma parameter. Please use SamplerTCD Euler A for the time being."

[JettHu]

From the link you provided: "NOTE: SamplerTCD is a WIP and currently just operates as DDIM with no gamma parameter. Please use SamplerTCD Euler A for the time being."

@RandomGitUser321 This PR does not require an additional TCD Scheduler. It uses sgm_uniform and simple in the BasicScheduler like lcm sampler.

The upper part of the screenshot below shows setting eta (gamma in paper) value through SamplerTCD. In the lower part, selecting tcd directly in KSampler will use the default 0.3 as gamma.

If necessary, I can adjust the PR and rename eta to gamma. Maybe I can add a gamma parameter to ModelSamplingDiscrete, or add a new node called ModelSamplingDiscreteTCD. cc @comfyanonymous

[RandomGitUser321]

From the link you provided: "NOTE: SamplerTCD is a WIP and currently just operates as DDIM with no gamma parameter. Please use SamplerTCD Euler A for the time being."

@RandomGitUser321 This PR does not require an additional TCD Scheduler. It uses sgm_uniform and simple in the BasicScheduler like lcm sampler.

The upper part of the screenshot below shows setting eta (gamma in paper) value through SamplerTCD. In the lower part, selecting tcd directly in KSampler will use the default 0.3 as gamma.

Oh I wasn't talking about your version, yours seems to work fine. I was talking about the dlf version. There are two versions of the comfyui implementation, yours https://github.com/JettHu/ComfyUI-TCD and his https://github.com/dfl/comfyui-tcd-scheduler

[JettHu]

Oh I wasn't talking about your version, yours seems to work fine. I was talking about the dlf version. There are two versions of the comfyui implementation, yours https://github.com/JettHu/ComfyUI-TCD and his https://github.com/dfl/comfyui-tcd-scheduler

Oh sorry, misunderstood

[dathide]

There's a low chance this gets merged. If you look at the commit history, Comfy hardly ever merges pull requests that aren't bug fixes. I'm not complaining, this is their repo that works well for me after all.

[JettHu]

There's a low chance this gets merged. If you look at the commit history, Comfy hardly ever merges pull requests that aren't bug fixes. I'm not complaining, this is their repo that works well for me after all.

That's a bit of a pity

[comfyanonymous]

This one is getting merged once I check that it matches the reference code.

[comfyanonymous]

official example:

this PR:

The quality of the images seems to be slightly worse than the official example

[MoonRide303]

@comfyanonymous It works pretty fine for me - via SamplerTCD EulerA with a bit higher gamma:

SDXL + TCD workflow, original prompt.json

[MoonRide303]

@mhh0318 @jabir-zheng @JettHu Could you guys please look into that difference mentioned by @comfyanonymous (between original diffusers and ComfyUI implementation), and double-check if sampler part is implemented properly?

I really like the output I get from this sampler, and I think it deserves to be inlcuded - just please make sure it's working as intended, so we wouldn't have to change it after users start using it.