[Code] Scheduler

Sigma 스케줄 코드 정리

샘플러는 매 step마다 사용할 sigma 값을 필요로 한다. 큰 sigma에서 시작해서 0에 가까운 값으로 내려가며 노이즈를 줄인다.

여기서 중요한 건 sigma의 의미보다, 코드에서 이 값들을 어떻게 만드는지다. 같은 20 step이라도 sigma를 어떤 간격으로 배치하느냐에 따라 결과가 달라진다.

이 코드에서는 네 가지 스케줄을 다룬다.

• linear: 모델의 기본 timestep을 그대로 사용
• exponential: log sigma 공간에서 균등 분할
• karras: 작은 sigma 구간에 step을 더 몰아주는 방식
• ays: 미리 정해진 timestep 목록을 사용

전체 코드: mi-song/image-generation
  ├ utils/schedule.py — 스케줄 함수
  └ components/samplers/base.py — 샘플러에서 스케줄 선택

---

1. 기본 sigma 만들기

가장 기본이 되는 값은 모델이 원래 쓰는 timestep에서 가져온다. 코드에서는 self.model.timesteps와 self.model.total_alphas를 사용한다.

def _native_sigmas(self) -> torch.Tensor: ts = self.model.timesteps.cpu().long() alphas = self.model.total_alphas[ts] sigmas = ((1 - alphas) / alphas).sqrt() return torch.cat([sigmas, sigmas.new_zeros([1])])

components/samplers/base.py L70-L74

여기서 하는 일은 단순하다.

• timesteps에서 사용할 timestep 목록을 가져온다.
• 각 timestep에 해당하는 alpha 값을 lookup한다.
• ((1 - alpha) / alpha).sqrt()로 sigma를 만든다.
• 마지막에 0을 하나 붙인다.

마지막 0은 샘플링 루프에서 마지막 도착점처럼 쓰인다. 그래서 길이는 step 수보다 하나 더 길어진다.

---

2. Linear

linear는 별도 함수가 없다. 위에서 만든 native sigma를 그대로 쓴다.

else: # "linear" return base

이 방식은 가장 기본 동작에 가깝다. 모델이 정한 timestep 분포를 그대로 따라가기 때문에 호환성이 좋다.

대신 step 수가 적을 때는 다른 스케줄보다 효율이 떨어질 수 있다. sigma 구간을 샘플링용으로 다시 최적화하지 않기 때문이다.

---

3. Exponential

exponential은 sigma를 log 공간에서 균등하게 나눈다.

def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'): sigmas = torch.linspace( sigma_max.log(), sigma_min.log(), n, device=device ).exp() return torch.cat([sigmas, sigmas.new_zeros([1])])

utils/schedule.py L49-L51

코드 흐름은 거의 한 줄이다.

• sigma_max.log()에서 시작한다.
• sigma_min.log()까지 n개로 나눈다.
• 마지막에 exp()로 원래 sigma 값으로 되돌린다.

이렇게 하면 sigma가 일정한 비율로 줄어든다. 큰 값에서는 크게 줄고, 작은 값에서는 더 촘촘하게 줄어드는 형태가 된다.

---

4. Karras

karras는 작은 sigma 구간에 step을 더 많이 배치하는 방식이다. 이미지의 후반 디테일을 잡는 구간에 계산을 더 쓰는 쪽에 가깝다.

def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'): ramp = torch.linspace(0, 1, n + 1, device=device)[:-1] min_inv_rho = sigma_min ** (1 / rho) max_inv_rho = sigma_max ** (1 / rho) sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return torch.cat([sigmas, sigmas.new_zeros([1])]).to(device)

utils/schedule.py L41-L46

코드에서 볼 부분은 세 군데다.

• ramp는 0에서 1로 가는 진행률이다.
• sigma_min ** (1 / rho), sigma_max ** (1 / rho)로 sigma를 다른 공간으로 옮긴다.
• 그 공간에서 보간한 뒤 다시 ** rho를 해서 sigma로 되돌린다.

rho가 곡선의 휘어짐을 정한다. 기본값은 7이다.

실제로 많이 쓰이는 이유는 간단하다. 같은 step 수에서 후반부 디테일이 더 안정적으로 나오는 경우가 많다.

---

5. AYS

ays는 앞의 두 방식처럼 수식으로 곡선을 만드는 방식이 아니다. 모델별로 정해둔 timestep 목록을 먼저 사용한다.

_AYS_TIMESTEPS = { "sdxl": [999, 845, 730, 587, 443, 310, 193, 116, 53, 13], "sd": [999, 850, 736, 645, 545, 455, 343, 233, 124, 24], }

utils/schedule.py L5-L8

SDXL과 SD 1.5의 목록이 따로 있다. 그래서 코드에서 먼저 모델 종류를 고른다.

def get_sigmas_ays(n, total_alphas, is_sdxl=True, device="cpu"): key = "sdxl" if is_sdxl else "sd" base_ts = torch.tensor(_AYS_TIMESTEPS[key], dtype=torch.long) base_alphas = total_alphas[base_ts] base_sigmas = ((1 - base_alphas) / base_alphas).sqrt().float() if n == len(base_ts): sigmas = base_sigmas else: t_base = np.linspace(0, 1, len(base_ts)) t_new = np.linspace(0, 1, n) log_interp = np.interp( t_new, t_base, base_sigmas.log().numpy() ) sigmas = torch.from_numpy(np.exp(log_interp)).float() return torch.cat([ sigmas.to(device), torch.zeros(1, device=device) ])

utils/schedule.py L11-L30

흐름은 이렇게 보면 된다.

• AYS timestep 목록을 가져온다.
• total_alphas에서 해당 timestep의 alpha를 찾는다.
• alpha를 sigma로 변환한다.
• 요청한 step 수가 10이면 그대로 쓴다.
• 10이 아니면 log sigma 공간에서 보간한다.

보간을 log sigma에서 하는 이유는 sigma 값의 범위가 꽤 크기 때문이다. 그냥 sigma 값으로 선형 보간하면 작은 sigma 구간이 뭉개지기 쉽다.

SDXL 여부 확인

AYS는 SDXL용 목록과 SD용 목록이 다르다. 이 코드에서는 tokenizer_2가 있는지로 SDXL 여부를 판단한다.

is_sdxl = hasattr(self.model, "tokenizer_2")

SDXL은 텍스트 인코더가 두 개라 두 번째 tokenizer를 가진다. 그래서 이 속성이 있으면 SDXL로 보고, 없으면 SD 계열로 처리한다.

---

6. 샘플러에서 스케줄 고르기

실제로 샘플러가 호출하는 쪽은 _get_sigmas()다. 여기서 문자열 설정에 따라 어떤 스케줄을 쓸지 나뉜다.

def _get_sigmas(self) -> torch.Tensor: base = self._native_sigmas() sigma_max, sigma_min = base[0], base[-2] n = len(self.model.timesteps) if self.schedule == "karras": return get_sigmas_karras( n, sigma_min, sigma_max, device=self.model.device ) elif self.schedule == "exponential": return get_sigmas_exponential( n, sigma_min, sigma_max, device=self.model.device ) elif self.schedule == "ays": is_sdxl = hasattr(self.model, "tokenizer_2") return get_sigmas_ays( n, self.model.total_alphas, is_sdxl=is_sdxl, device=self.model.device ) else: return base

components/samplers/base.py L76-L90

여기서 base[-2]를 쓰는 이유는 마지막 값이 항상 0이기 때문이다. 실제 최소 sigma는 0 바로 앞에 있는 값이다.

스케줄별로 필요한 인자도 조금 다르다.

스케줄	필요한 값	코드 기준 차이
linear	없음	base 그대로 반환
exponential	sigma_min, sigma_max	log sigma 공간에서 보간
karras	sigma_min, sigma_max	rho를 이용해 곡선 보간
ays	total_alphas	정해진 timestep을 alpha 테이블에서 직접 조회

---

정리

코드 기준으로 보면 sigma 스케줄은 결국 N+1 길이의 sigma 텐서를 만드는 함수다.

스케줄	짧게 보면
linear	모델 timestep을 그대로 sigma로 바꾼다.
exponential	log sigma를 균등하게 나눈다.
karras	작은 sigma 쪽에 step을 더 배치한다.
ays	모델별로 정해둔 timestep 목록을 사용한다.

샘플링 루프 입장에서는 어떤 스케줄이든 같은 형식의 sigma 배열만 받으면 된다. 차이는 그 배열을 만드는 방식에 있다.