Diffusers documentation
TCDScheduler
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
< source >( num_train_timesteps: int = 1000beta_start: float = 0.00085beta_end: float = 0.012beta_schedule: str = 'scaled_linear'trained_betas: numpy.ndarray | list[float] | None = Noneoriginal_inference_steps: int = 50clip_sample: bool = Falseclip_sample_range: float = 1.0set_alpha_to_one: bool = Truesteps_offset: int = 0prediction_type: str = 'epsilon'thresholding: bool = Falsedynamic_thresholding_ratio: float = 0.995sample_max_value: float = 1.0timestep_spacing: str = 'leading'timestep_scaling: float = 10.0rescale_betas_zero_snr: bool = False )
매개변수
- num_train_timesteps (
int, 기본값1000) — 모델을 학습할 때 사용할 diffusion 단계 수입니다. - 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 곱 값을 사용합니다. 마지막 단계에는 이전 alpha가 없습니다. 이 옵션이True이면 이전 alpha 곱 값은1로 고정되고, 그렇지 않으면 0단계의 alpha 값을 사용합니다. - steps_offset (
int, 기본값0) — 일부 모델 계열에서 요구하는 대로 추론 단계에 추가되는 오프셋입니다. - prediction_type (
str, defaults to"epsilon") — Prediction type of the scheduler function; can beepsilon(predicts the noise of the diffusion process),sample(directly predicts the noisy sample) orv_prediction` (see section 2.4 of Imagen Video paper). - thresholding (
bool, 기본값False) — “동적 임계값 설정” 방법을 사용할지 여부입니다. Stable Diffusion과 같은 잠재 공간 확산 모델에는 적합하지 않습니다. - dynamic_thresholding_ratio (
float, 기본값0.995) — 동적 임계값 설정 방법의 비율입니다.thresholding=True일 때만 유효합니다. - sample_max_value (
float, 기본값1.0) — 동적 임계값 설정을 위한 임계값입니다.thresholding=True일 때만 유효합니다. - timestep_spacing (
str, 기본값"leading") — timestep을 스케일링하는 방법입니다. 자세한 내용은 공통 확산 노이즈 스케줄 및 샘플 단계는 결함이 있다의 표 2를 참조하세요. - timestep_scaling (
float, 기본값10.0) — consistency model의 경계 조건인c_skip과c_out을 계산할 때 timestep에 곱해지는 계수입니다. 이 값을 증가시키면 근사 오차가 감소합니다(기본값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__ 함수에 전달되는 모든 config 속성을 저장하는 역할을 합니다. 이러한 속성은 scheduler.config.num_train_timesteps을 통해 접근할 수 있습니다. SchedulerMixin은 SchedulerMixin.save_pretrained() 및 from_pretrained() 함수를 통해 일반적인 로딩 및 저장 기능을 제공합니다.
add_noise
< source >( original_samples: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor
각 timestep의 노이즈 크기에 따라 원래 샘플에 노이즈를 추가합니다(이것은 forward diffusion process입니다).
get_velocity
< source >( sample: Tensornoise: Tensortimesteps: IntTensor ) → torch.Tensor
속도 공식에 따라 샘플과 노이즈로부터 속도 예측을 계산합니다.
index_for_timestep
< source >( timestep: typing.Union[float, torch.Tensor]schedule_timesteps: typing.Optional[torch.Tensor] = None ) → int
…
이전 시간 단계
< source >( timestep: int ) → int or torch.Tensor
확산 체인의 이전 시간 단계를 계산합니다.
scale_model_input
< source >( sample: Tensortimestep: typing.Union[int, torch.Tensor, NoneType] = None ) → torch.Tensor
현재 timestep에 따라 denoising 모델 입력을 scale해야 하는 scheduler와의 호환성을 보장합니다.
set_begin_index
< source >( begin_index: int = 0 )
… 스케줄러의 시작 인덱스를 설정합니다. 이 함수는 추론 전에 파이프라인에서 실행해야 합니다.
set_timesteps
< source >( num_inference_steps: int | None = Nonedevice: typing.Union[str, torch.device, NoneType] = Noneoriginal_inference_steps: int | None = Nonetimesteps: list[int] | None = Nonestrength: float = 1.0 )
매개변수
- num_inference_steps (
int, 선택 사항) — 학습된 모델로 샘플을 생성할 때 사용되는 diffusion 단계 수입니다. 이 값을 사용하는 경우timesteps은None이어야 합니다. - device (
str또는torch.device, 선택 사항) — timesteps를 이동할 대상 device입니다.None인 경우 timesteps는 이동되지 않습니다. - original_inference_steps (
int, 선택 사항) — 선형 간격의 타임스텝 스케줄을 생성하는 데 사용되는 원래 추론 단계 수입니다(표준diffusers구현과 다름). 이후 이 스케줄에서 인덱스 기준으로 균등하게 간격을 두고num_inference_steps개의 타임스텝을 선택하여 최종 타임스텝 스케줄로 사용합니다. 설정되지 않으면original_inference_steps속성 값이 기본값으로 사용됩니다. - timesteps (
list[int], 선택 사항) — timesteps 사이의 임의의 간격을 지원하기 위해 사용되는 사용자 지정 timesteps입니다.None인 경우, 학습/증류 timestep 스케줄에서 timesteps 사이에 동일한 간격을 두는 기본 timestep 간격 전략이 사용됩니다.timesteps가 전달되면num_inference_steps는None이어야 합니다. - strength (
float, 기본값1.0) — img2img, inpaint 등을 사용할 때 추론에 사용되는 타임스텝 수를 결정하는 데 사용됩니다.
추론 실행 전에 실행할 diffusion chain에 사용할 이산 timesteps를 설정합니다.
단계
< source >( model_output: Tensortimestep: typing.Union[int, torch.Tensor]sample: Tensoreta: float = 0.3generator: typing.Optional[torch.Generator] = Nonereturn_dict: bool = True ) → TCDSchedulerOutput or tuple[torch.Tensor, torch.Tensor]
매개변수
- model_output (
torch.Tensor) — 학습된 diffusion 모델의 직접적인 출력입니다. - 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) —tuple대신 TCDSchedulerOutput를 반환할지 여부.
반환값
TCDSchedulerOutput 또는 tuple[torch.Tensor, torch.Tensor]
return_dict가 True인 경우 TCDSchedulerOutput이 반환되고, 그렇지 않으면 이전 샘플과 예측된 노이즈 샘플을 포함하는 튜플이 반환됩니다.
이전 시간 단계의 샘플을 SDE를 역방향으로 적용하여 예측합니다. 이 함수는 학습된 모델의 출력(대부분 예측된 노이즈)으로부터 확산 과정을 전파합니다.
TCDSchedulerOutput
class diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput
< source >( prev_sample: Tensorpred_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, 이미지의 경우(batch_size, num_channels, height, width)형태, 선택 사항) — 현재 timestep의 모델 출력에 기반한 예측된 노이즈 샘플(x_{s}).
스케줄러의 step 함수 출력을 위한 출력 클래스입니다.