Fix static type errors with SCHEDULER_NAME_VALUES. And, avoid bi-directional cross-directory imports, which contribute to circular import issues.

This commit is contained in:
Ryan Dick
2024-07-03 11:13:16 -04:00
committed by Kent Keirsey
parent 3a24d70279
commit 35f8781ea2
8 changed files with 49 additions and 11 deletions

View File

@ -1,3 +1,5 @@
from typing import Any, Literal, Type
from diffusers import (
DDIMScheduler,
DDPMScheduler,
@ -16,8 +18,36 @@ from diffusers import (
TCDScheduler,
UniPCMultistepScheduler,
)
from diffusers.schedulers.scheduling_utils import SchedulerMixin
SCHEDULER_MAP = {
SCHEDULER_NAME_VALUES = Literal[
"ddim",
"ddpm",
"deis",
"lms",
"lms_k",
"pndm",
"heun",
"heun_k",
"euler",
"euler_k",
"euler_a",
"kdpm_2",
"kdpm_2_a",
"dpmpp_2s",
"dpmpp_2s_k",
"dpmpp_2m",
"dpmpp_2m_k",
"dpmpp_2m_sde",
"dpmpp_2m_sde_k",
"dpmpp_sde",
"dpmpp_sde_k",
"unipc",
"lcm",
"tcd",
]
SCHEDULER_MAP: dict[SCHEDULER_NAME_VALUES, tuple[Type[SchedulerMixin], dict[str, Any]]] = {
"ddim": (DDIMScheduler, {}),
"ddpm": (DDPMScheduler, {}),
"deis": (DEISMultistepScheduler, {}),