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

UPDATE: normal EulerA seems to do pretty well (maybe even a bit better?) with TCD LoRA, too:
workflow SDXL + TCD, EulerA.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?

UPDATE: at the moment I am not sure how much of myself being impressed with TCD was based on TCD sampler / scheduler, or it was just TCD LoRA. I am getting nice results from this LoRA when simply using it with sgm_uniform scheduler and EulerA sampler.

[JettHu]

UPDATE: at the moment I am not sure how much of myself being impressed with TCD was based on TCD sampler / scheduler, or it was just TCD LoRA. I am getting nice results from this LoRA when simply using it with sgm_uniform scheduler and EulerA sampler.

Okay, let me take a look at this situation.

[JettHu]

official example:

this PR:

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

Are the results of the official example generated by diffusers?
In addition, can you share the parameters or workflow of the PR result?

[JettHu]

By the way, my results and workflow.

tcd_pr_workflow.json

[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?

UPDATE: at the moment I am not sure how much of myself being impressed with TCD was based on TCD sampler / scheduler, or it was just TCD LoRA. I am getting nice results from this LoRA when simply using it with sgm_uniform scheduler and EulerA sampler.

After my test EulerA sampler can be used on TCD-LoRA, and EulerA sampler can also be used on LCM-LoRA. It gives decent results, but for better work it requires a special sampler. Seems to have helped with Hyper-SD.

And I re-read the paper as well as the tcd sampler implementation of diffusers. I feel there is nothing wrong. It would be better if @mhh0318 and @jabir-zheng could help take a look.

[mhh0318]

Nice work! I have checked it. Thanks for your contribution. I feel this Euler based sampler can also be used in SD3 version.

[comfyanonymous]

I tried your exact workflow and branch but I cannot reproduce your result, this is what I get:

What hardware do you have? Do you have any other modifications to ComfyUI?

[JettHu]

Tesla T4 8G VRAM, Driver Version: 515.65.01 CUDA Version: 11.7

Server is started with --disable-cuda-malloc

# pip freeze | grep torch
torch==2.1.1+cu118
torchsde==0.2.6
torchvision==0.16.1+cu118
# I've not install xformers

I'll try another device.

[bigmover]

Nice work! I have checked it. Thanks for your contribution. I feel this Euler based sampler can also be used in SD3 version.

Wether Euler sampler is good enough to support TCD ? Would you mind to give some advice or help to release an official support for tcd?

[dfl]

@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?

UPDATE: at the moment I am not sure how much of myself being impressed with TCD was based on TCD sampler / scheduler, or it was just TCD LoRA. I am getting nice results from this LoRA when simply using it with sgm_uniform scheduler and EulerA sampler.

I think there is may be some confusion about eta = 1 - gamma? TCD diffusers repo says when eta=0, it is deterministic sampling (which AFAIK should be identical to EulerA ?). I think this is one of the nicer features of TCD, that you can subtly adjust the detail by modifying this parameter -- it's like a crossfade between Euler and DDPM. For example, this type of detail is nice for features like wood grain or hair.