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
매개변수
- timestep (
float또는torch.Tensor) — 스케줄에서 찾을 timestep 값입니다. - schedule_timesteps (
torch.Tensor, 선택 사항) — 검색할 timestep schedule입니다.None인 경우self.timesteps를 사용합니다.
반환값
int
The index of the timestep in the schedule. For the very first step, returns the second index if multiple matches exist to avoid skipping a sigma when starting mid-schedule (e.g., for image-to-image).
…
이전 시간 단계
< 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 함수 출력을 위한 출력 클래스입니다.