diffusers/main/ko/api/schedulers/tcd.md

TCDScheduler

Jianbin Zheng, Minghui Hu, Zhongyi Fan, Chaoyue Wang, Changxing Ding, Dacheng Tao, Tat-Jen Cham의 Trajectory Consistency Distillation은 소수의 단계로 양질의 샘플을 생성할 수 있는 Strategic Stochastic Sampling(Algorithm 4)을 소개했습니다. Consistency Models의 multistep scheduler(Algorithm 1)의 고급 반복 버전으로 구분되는 Strategic Stochastic Sampling은 trajectory consistency function에 특화되어 있습니다.

해당 논문의 초록은 다음과 같습니다:

Latent Consistency Model (LCM)은 Consistency Model을 latent space로 확장하고, guided consistency distillation 기법을 활용하여 text-to-image 합성 가속화에서 인상적인 성능을 달성합니다. 그러나 LCM이 선명함과 세밀한 디테일을 모두 갖춘 이미지를 생성하는 데 어려움을 겪는다는 점을 관찰했습니다. 이 한계를 해결하기 위해, 우리는 먼저 근본적인 원인을 심층적으로 분석하고 설명합니다. 우리의 조사에 따르면, 주요 문제는 세 가지 서로 다른 영역의 오류에서 비롯됩니다. 이에 따라, trajectory consistency function과 strategic stochastic sampling을 포함하는 Trajectory Consistency Distillation (TCD)을 도입합니다. trajectory consistency function은 self-consistency boundary condition의 범위를 넓혀 distillation 오류를 줄이고, TCD가 Probability Flow ODE의 전체 trajectory를 정확하게 추적할 수 있는 능력을 부여합니다. 또한, strategic stochastic sampling은 multi-step consistency sampling에 내재된 누적 오류를 우회하도록 특별히 설계되었으며, TCD 모델을 보완하도록 정밀하게 조정되었습니다. 실험 결과, TCD는 낮은 NFE에서 이미지 품질을 크게 향상시킬 뿐만 아니라, 높은 NFE에서는 teacher 모델보다 더 상세한 결과를 산출하는 것으로 나타났습니다.

원본 코드베이스는 jabir-zheng/TCD에서 찾을 수 있습니다.

TCDScheduler[[diffusers.TCDScheduler]]

diffusers.TCDScheduler[[diffusers.TCDScheduler]]

diffusers.TCDScheduler(num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: str = 'scaled_linear', trained_betas: numpy.ndarray | list[float] | None = None, original_inference_steps: int = 50, clip_sample: bool = False, clip_sample_range: float = 1.0, set_alpha_to_one: bool = True, steps_offset: int = 0, prediction_type: str = 'epsilon', thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, timestep_spacing: str = 'leading', timestep_scaling: float = 10.0, rescale_betas_zero_snr: bool = False)

소스

매개변수:

num_train_timesteps (int, 기본값 1000) : 모델을 학습하기 위한 diffusion step 수입니다.

beta_start (float, 기본값 0.00085) : 추론의 시작 beta 값입니다.

beta_end (float, 기본값 0.012) : 최종 beta 값.

beta_schedule (str, 기본값 "scaled_linear") : 베타 스케줄로, 모델의 스테핑을 위한 베타 범위에서 베타 시퀀스로의 매핑입니다. linear, scaled_linear, 또는 squaredcos_cap_v2 중 하나를 선택합니다.

trained_betas (np.ndarray 또는 list[float], 선택 사항) : beta_start와 beta_end를 우회하여 베타 배열을 생성자에 직접 전달합니다.

original_inference_steps (int, 기본값 50) : 선형 간격의 timestep 스케줄을 생성하는 데 사용되는 기본 추론 단계 수입니다. 이 스케줄에서 최종적으로 num_inference_steps개의 균등 간격 timestep을 추출하여 최종 timestep 스케줄을 구성합니다.

clip_sample (bool, 기본값 False) : 수치적 안정성을 위해 예측된 샘플을 클리핑합니다.

clip_sample_range (float, 기본값 1.0) : 샘플 클리핑의 최대 크기입니다. clip_sample=True일 때만 유효합니다.

set_alpha_to_one (bool, 기본값 True) : 각 diffusion 단계는 해당 단계와 이전 단계의 alphas product 값을 사용합니다. 마지막 단계에는 이전 alpha가 없습니다. 이 옵션이 True인 경우 이전 alpha product는 1로 고정되고, 그렇지 않으면 step 0의 alpha 값을 사용합니다.

steps_offset (int, 기본값 0) : 일부 모델 계열에서 요구하는 추론 단계에 추가되는 오프셋입니다.

prediction_type (str, defaults to "epsilon") : Prediction type of the scheduler function; can be epsilon (predicts the noise of the diffusion process), sample (directly predicts the noisy sample) or v_prediction` (see section 2.4 of Imagen Video paper).

thresholding (bool, 기본값은 False) : "dynamic thresholding" 방법을 사용할지 여부입니다. Stable Diffusion과 같은 latent-space diffusion 모델에는 적합하지 않습니다.

dynamic_thresholding_ratio (float, 기본값 0.995) : 동적 임계값 설정 방법의 비율입니다. thresholding=True일 때만 유효합니다.

sample_max_value (float, 기본값 1.0) : 동적 임계값 설정의 임계값입니다. thresholding=True일 때만 유효합니다.

timestep_spacing (str, 기본값 "leading") : timesteps를 스케일링하는 방식입니다. 자세한 내용은 일반적인 디퓨전 노이즈 스케줄과 샘플 스테프는 결함이 있다의 Table 2를 참조하세요.

timestep_scaling (float, 기본값 10.0) : consistency model의 경계 조건 c_skip과 c_out을 계산할 때 timesteps에 곱해지는 계수입니다. 이 값을 증가시키면 근사 오차가 감소합니다(기본값 10.0에서의 근사 오차는 이미 상당히 작습니다).

rescale_betas_zero_snr (bool, 기본값 False) : 베타를 재조정하여 최종 SNR이 0이 되도록 할지 여부입니다. 이를 통해 모델은 중간 밝기의 샘플로 제한되는 대신 매우 밝고 어두운 샘플을 생성할 수 있습니다. --offset_noise와 느슨하게 관련이 있습니다.

TCDScheduler은 Trajectory Consistency Distillation 논문에서 도입한 Strategic Stochastic Sampling을 포함하며, 원래의 Multistep Consistency Sampling을 확장하여 제한 없는 궤적 탐색을 가능하게 합니다.

이 코드는 TCD의 공식 저장소(https://github.com/jabir-zheng/TCD)를 기반으로 합니다.

이 모델은 SchedulerMixin과 ConfigMixin을 상속합니다. ~ConfigMixin는 num_train_timesteps와 같은 스케줄러의 __init__ 함수에 전달되는 모든 설정 속성을 저장하는 역할을 합니다. 이러한 속성은 scheduler.config.num_train_timesteps을 통해 접근할 수 있습니다. SchedulerMixin은 SchedulerMixin.save_pretrained() 및 from_pretrained() 함수를 통해 일반적인 로딩 및 저장 기능을 제공합니다.

add_noise[[diffusers.TCDScheduler.add_noise]]

add_noise(original_samples: Tensor, noise: Tensor, timesteps: IntTensor)

소스

매개변수:

original_samples (torch.Tensor) : 노이즈가 추가될 원래 샘플.

noise (torch.Tensor) : 샘플에 추가할 노이즈.

timesteps (torch.IntTensor) : 각 샘플의 노이즈 수준을 나타내는 timesteps.

반환값: torch.Tensor

노이즈가 포함된 샘플.

각 timestep의 노이즈 크기에 따라 원래 샘플에 노이즈를 추가합니다(이것은 forward diffusion process입니다).

get_velocity[[diffusers.TCDScheduler.get_velocity]]

get_velocity(sample: Tensor, noise: Tensor, timesteps: IntTensor)

소스

매개변수:

sample (torch.Tensor) : 입력 샘플.

noise (torch.Tensor) : 노이즈 텐서.

timesteps (torch.IntTensor) : 속도 계산에 사용되는 timesteps입니다.

반환값: torch.Tensor

계산된 속도.

속도 공식에 따라 샘플과 노이즈로부터 속도 예측을 계산합니다.

index_for_timestep[[diffusers.TCDScheduler.index_for_timestep]]

index_for_timestep(timestep: typing.Union[float, torch.Tensor], schedule_timesteps: typing.Optional[torch.Tensor] = None)

소스

매개변수:

timestep (float 또는 torch.Tensor) : 스케줄에서 찾을 타임스텝 값.

schedule_timesteps (torch.Tensor, 선택 사항) : 검색할 timestep schedule입니다. None인 경우 self.timesteps를 사용합니다.

반환값: int

스케줄에서 timestep의 인덱스입니다. 가장 첫 번째 단계의 경우, 스케줄 중간에서 시작할 때(예: image-to-image) sigma를 건너뛰지 않도록 여러 매치가 존재하면 두 번째 인덱스를 반환합니다.

…

previous_timestep[[diffusers.TCDScheduler.previous_timestep]]

previous_timestep(timestep: int)

소스

매개변수:

timestep (int) : 현재 타임스텝.

반환값: int 또는 torch.Tensor

이전 시간 단계(timestep).

확산 체인의 이전 시간 단계를 계산합니다.

scale_model_input[[diffusers.TCDScheduler.scale_model_input]]

scale_model_input(sample: Tensor, timestep: typing.Union[int, torch.Tensor, NoneType] = None)

소스

매개변수:

sample (torch.Tensor) : 입력 샘플.

timestep (int 또는 torch.Tensor, 선택 사항) : 확산 체인(diffusion chain)의 현재 시간 단계(timestep).

반환값: torch.Tensor

스케일된 입력 샘플.

현재 타임스텝에 따라 디노이징 모델 입력을 스케일링해야 하는 스케줄러와의 호환성을 보장합니다.

set_begin_index[[diffusers.TCDScheduler.set_begin_index]]

set_begin_index(begin_index: int = 0)

소스

매개변수:

begin_index (int, 기본값 0) : 스케줄러의 시작 인덱스입니다.

… 스케줄러의 시작 인덱스를 설정합니다. 이 함수는 추론 전에 파이프라인에서 실행해야 합니다.

set_timesteps[[diffusers.TCDScheduler.set_timesteps]]

set_timesteps(num_inference_steps: int | None = None, device: typing.Union[str, torch.device, NoneType] = None, original_inference_steps: int | None = None, timesteps: list[int] | None = None, strength: float = 1.0)

소스

매개변수:

num_inference_steps (int, 선택 사항) : 사전 학습된 모델로 샘플을 생성할 때 사용되는 확산 단계 수입니다. 이 값을 사용하는 경우 timesteps는 None이어야 합니다.

device (str 또는 torch.device, 선택 사항) : timesteps를 이동할 대상 device입니다. None인 경우 timesteps는 이동되지 않습니다.

original_inference_steps (int, 선택 사항) : 선형 간격의 timestep 스케줄을 생성하는 데 사용되는 원래 추론 단계 수입니다(표준 diffusers 구현과 다름). 이후 이 스케줄에서 인덱스 기준으로 균등하게 간격을 두고 num_inference_steps개의 timesteps를 선택하여 최종 timestep 스케줄로 사용합니다. 설정되지 않으면 original_inference_steps 속성 값이 기본값으로 사용됩니다.

timesteps (list[int], 선택 사항) : timestep 간 임의의 간격을 지원하기 위해 사용되는 사용자 지정 timestep입니다. None인 경우, 학습/증류 timestep 스케줄에서 timestep 간 균등 간격이라는 기본 timestep 간격 전략이 사용됩니다. timesteps가 전달되면 num_inference_steps는 None이어야 합니다.

strength (float, 기본값 1.0) : img2img, inpaint 등을 사용할 때 추론에 사용되는 timestep 수를 결정하는 데 사용됩니다.

추론 실행 전에 실행할 diffusion chain에 사용할 이산 timesteps를 설정합니다.

step[[diffusers.TCDScheduler.step]]

step(model_output: Tensor, timestep: typing.Union[int, torch.Tensor], sample: Tensor, eta: float = 0.3, generator: typing.Optional[torch.Generator] = None, return_dict: bool = True)

소스

매개변수:

model_output (torch.Tensor) : 학습된 diffusion model의 직접적인 출력입니다.

timestep (int 또는 torch.Tensor) : 확산 체인(diffusion chain)의 현재 이산 시간 단계(discrete timestep).

sample (torch.Tensor) : 확산 과정에서 생성된 샘플의 현재 인스턴스.

eta (float, 기본값 0.3) : 각 단계의 확률성을 제어하는 데 사용되는 확률적 파라미터(논문에서는 gamma로 지칭). eta = 0이면 결정론적 샘플링을, eta = 1이면 완전한 확률적 샘플링을 나타냅니다.

generator (torch.Generator, 선택 사항) : 난수 생성기입니다.

return_dict (bool, 기본값 True) : TCDSchedulerOutput 또는 tuple를 반환할지 여부.

반환값: TCDSchedulerOutput 또는 tuple[torch.Tensor, torch.Tensor]

return_dict가 True인 경우 TCDSchedulerOutput이 반환되고, 그렇지 않으면 이전 샘플과 예측된 노이즈 샘플을 포함하는 튜플이 반환됩니다.

이전 시간 단계의 샘플을 SDE를 역방향으로 적용하여 예측합니다. 이 함수는 학습된 모델의 출력(대부분 예측된 노이즈)으로부터 확산 과정을 전파합니다.

TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput(prev_sample: Tensor, pred_noised_sample: typing.Optional[torch.Tensor] = None)

소스

매개변수:

prev_sample (torch.Tensor, 이미지인 경우 shape는 (batch_size, num_channels, height, width)) : 이전 timestep의 계산된 샘플 (x_{t-1}). prev_sample은 denoising 루프에서 다음 모델 입력으로 사용해야 합니다.

pred_noised_sample (torch.Tensor, 이미지인 경우 shape (batch_size, num_channels, height, width), optional) : 현재 timestep의 모델 출력에 기반한 예측된 노이즈 샘플 (x_{s}).

스케줄러의 step 함수 출력을 위한 출력 클래스입니다.