diff options
Diffstat (limited to 'src/mcp_sd15')
| -rw-r--r-- | src/mcp_sd15/config.py | 6 | ||||
| -rw-r--r-- | src/mcp_sd15/pipeline_factory.py | 33 |
2 files changed, 27 insertions, 12 deletions
diff --git a/src/mcp_sd15/config.py b/src/mcp_sd15/config.py index 5bbe225..1a5a71d 100644 --- a/src/mcp_sd15/config.py +++ b/src/mcp_sd15/config.py @@ -11,7 +11,7 @@ from typing import Any from .errors import ConfigError VALID_SIZES = frozenset({"square_512", "landscape_512", "portrait_512"}) -VALID_SCHEDULERS = frozenset({"dpmpp_2m", "ddim", "euler", "euler_a"}) +VALID_SCHEDULERS = frozenset({"dpmpp_2m_karras", "dpmpp_2m_sde_karras", "euler_a", "ddim", "lcm"}) @dataclass(frozen=True) @@ -41,7 +41,7 @@ class ModelConfig: class GenerationConfig: steps: int = 25 cfg_scale: float = 7.0 - scheduler: str = "dpmpp_2m" + scheduler: str = "dpmpp_2m_karras" default_size: str = "square_512" @@ -189,7 +189,7 @@ def load_config(path: str | Path, *, validate_paths: bool = True) -> AppConfig: cfg_scale = _number(generation_raw, "cfg_scale", 7.0) if not 0 <= cfg_scale <= 30: raise ConfigError("generation.cfg_scale must be between 0 and 30") - scheduler = _string(generation_raw, "scheduler") or "dpmpp_2m" + scheduler = _string(generation_raw, "scheduler") or "dpmpp_2m_karras" if scheduler not in VALID_SCHEDULERS: raise ConfigError(f"unsupported scheduler: {scheduler}") default_size = _string(generation_raw, "default_size") or "square_512" diff --git a/src/mcp_sd15/pipeline_factory.py b/src/mcp_sd15/pipeline_factory.py index 12d95ba..b0525ea 100644 --- a/src/mcp_sd15/pipeline_factory.py +++ b/src/mcp_sd15/pipeline_factory.py @@ -64,7 +64,7 @@ class PipelineFactory: adapter_weights.append(item.weight) if adapter_names: pipeline.set_adapters(adapter_names, adapter_weights=adapter_weights) - pipeline.scheduler = self._scheduler(pipeline.scheduler, diffusers) + pipeline.scheduler = self._create_scheduler(pipeline.scheduler, diffusers) pipeline.set_progress_bar_config(disable=True) return pipeline.to("cuda") except ConfigError: @@ -130,14 +130,29 @@ class PipelineFactory: pipeline.set_progress_bar_config(disable=True) return pipeline.to("cuda") - def _scheduler(self, scheduler: Any, diffusers: Any) -> Any: + def _create_scheduler(self, current: Any, diffusers: Any) -> Any: + name = self.config.generation.scheduler + if name == "dpmpp_2m_karras": + return diffusers.DPMSolverMultistepScheduler.from_config( + current.config, + algorithm_type="dpmsolver++", + solver_order=2, + use_karras_sigmas=True, + ) + if name == "dpmpp_2m_sde_karras": + return diffusers.DPMSolverMultistepScheduler.from_config( + current.config, + algorithm_type="sde-dpmsolver++", + solver_order=2, + use_karras_sigmas=True, + ) classes = { - "dpmpp_2m": diffusers.DPMSolverMultistepScheduler, - "ddim": diffusers.DDIMScheduler, - "euler": diffusers.EulerDiscreteScheduler, "euler_a": diffusers.EulerAncestralDiscreteScheduler, + "ddim": diffusers.DDIMScheduler, + "lcm": diffusers.LCMScheduler, } - options = dict(scheduler.config) - if self.config.generation.scheduler == "dpmpp_2m": - options.update(algorithm_type="dpmsolver++", solver_order=2) - return classes[self.config.generation.scheduler].from_config(options) + try: + scheduler_class = classes[name] + except KeyError as exc: + raise ConfigError(f"unsupported scheduler: {name}") from exc + return scheduler_class.from_config(current.config) |