Add more sigma types, clean up config with mixins, rebuild frontend

core/config/config.py

@@ -8,10 +8,28 @@
 from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
 
 from core.config.samplers.sampler_config import SamplerConfig
+from core.types import SigmaType
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class BaseDiffusionMixin:
+    width: int = 512
+    height: int = 512
+    batch_count: int = 1
+    batch_size: int = 1
+    seed: int = -1
+    cfg_scale: int = 7
+    steps: int = 40
+    prompt: str = ""
+    negative_prompt: str = ""
+    sampler: Union[
+        int, str
+    ] = KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler.value
+    sigmas: SigmaType = ""
+
+
 @dataclass
 class QuantDict:
     vae_decoder: Optional[bool] = None
@@ -21,80 +39,32 @@ class QuantDict:
 
 
 @dataclass
-class Txt2ImgConfig:
+class Txt2ImgConfig(BaseDiffusionMixin):
     "Configuration for the text to image pipeline"
 
-    width: int = 512
-    height: int = 512
-    seed: int = -1
-    cfg_scale: int = 7
-    sampler: Union[
-        int, str
-    ] = KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler.value
-    prompt: str = ""
-    negative_prompt: str = ""
-    steps: int = 40
-    batch_count: int = 1
-    batch_size: int = 1
     self_attention_scale: float = 0.0
 
 
 @dataclass
-class Img2ImgConfig:
+class Img2ImgConfig(BaseDiffusionMixin):
     "Configuration for the image to image pipeline"
 
-    width: int = 512
-    height: int = 512
-    seed: int = -1
-    cfg_scale: int = 7
-    sampler: Union[
-        int, str
-    ] = KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler.value
-    prompt: str = ""
-    negative_prompt: str = ""
-    steps: int = 40
-    batch_count: int = 1
-    batch_size: int = 1
     resize_method: int = 0
     denoising_strength: float = 0.6
     self_attention_scale: float = 0.0
 
 
 @dataclass
-class InpaintingConfig:
+class InpaintingConfig(BaseDiffusionMixin):
     "Configuration for the inpainting pipeline"
 
-    prompt: str = ""
-    negative_prompt: str = ""
-    width: int = 512
-    height: int = 512
-    steps: int = 40
-    cfg_scale: int = 7
-    seed: int = -1
-    batch_count: int = 1
-    batch_size: int = 1
-    sampler: Union[
-        int, str
-    ] = KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler.value
     self_attention_scale: float = 0.0
 
 
 @dataclass
-class ControlNetConfig:
+class ControlNetConfig(BaseDiffusionMixin):
     "Configuration for the inpainting pipeline"
 
-    prompt: str = ""
-    negative_prompt: str = ""
-    width: int = 512
-    height: int = 512
-    seed: int = -1
-    cfg_scale: int = 7
-    steps: int = 40
-    batch_count: int = 1
-    batch_size: int = 1
-    sampler: Union[
-        int, str
-    ] = KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler.value
     controlnet: str = "lllyasviel/sd-controlnet-canny"
     controlnet_conditioning_scale: float = 1.0
     detection_resolution: int = 512

core/inference/ait/aitemplate.py

@@ -25,6 +25,7 @@
     ControlNetQueueEntry,
     Img2ImgQueueEntry,
     Job,
+    SigmaType,
     Txt2ImgQueueEntry,
 )
 from core.utils import convert_images_to_base64_grid, convert_to_image, resize
@@ -239,7 +240,7 @@ def manage_optional_components(
     def create_pipe(
         self,
         controlnet: str = "",
-        scheduler: Optional[Tuple[Any, bool]] = None,
+        scheduler: Optional[Tuple[Any, SigmaType]] = None,
         sampler_settings: Optional[dict] = None,
     ) -> "StableDiffusionAITPipeline":
         "Centralized way to create new pipelines."
@@ -266,7 +267,7 @@ def create_pipe(
             change_scheduler(
                 model=pipe,
                 scheduler=scheduler[0],
-                use_karras_sigmas=scheduler[1],
+                sigma_type=scheduler[1],
                 sampler_settings=sampler_settings,
             )
         return pipe
@@ -293,7 +294,7 @@ def txt2img(
     ) -> List[Image.Image]:
         "Generates images from text"
         pipe = self.create_pipe(
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 
@@ -383,7 +384,7 @@ def img2img(
     ) -> List[Image.Image]:
         "Generates images from images"
         pipe = self.create_pipe(
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 
@@ -444,7 +445,7 @@ def controlnet2img(
         "Generates images from images"
         pipe = self.create_pipe(
             controlnet=job.data.controlnet,
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 

core/inference/pytorch/pytorch.py

@@ -39,6 +39,7 @@
     Img2ImgQueueEntry,
     InpaintQueueEntry,
     Job,
+    SigmaType,
     Txt2ImgQueueEntry,
 )
 from core.utils import convert_images_to_base64_grid, convert_to_image, resize
@@ -264,7 +265,7 @@ def manage_optional_components(
     def create_pipe(
         self,
         controlnet: Optional[str] = "",
-        scheduler: Optional[Tuple[Any, bool]] = None,
+        scheduler: Optional[Tuple[Any, SigmaType]] = None,
         sampler_settings: Optional[dict] = None,
     ) -> StableDiffusionLongPromptWeightingPipeline:
         "Create a pipeline -- useful for reducing backend clutter."
@@ -284,7 +285,7 @@ def create_pipe(
             change_scheduler(
                 model=pipe,
                 scheduler=scheduler[0],  # type: ignore
-                use_karras_sigmas=scheduler[1],
+                sigma_type=scheduler[1],
                 sampler_settings=sampler_settings,
             )
 
@@ -294,7 +295,7 @@ def txt2img(self, job: Txt2ImgQueueEntry) -> List[Image.Image]:
         "Generate an image from a prompt"
 
         pipe = self.create_pipe(
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 
@@ -379,7 +380,7 @@ def img2img(self, job: Img2ImgQueueEntry) -> List[Image.Image]:
         "Generate an image from an image"
 
         pipe = self.create_pipe(
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 
@@ -440,7 +441,7 @@ def inpaint(self, job: InpaintQueueEntry) -> List[Image.Image]:
         "Generate an image from an image"
 
         pipe = self.create_pipe(
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 
@@ -512,7 +513,7 @@ def controlnet2img(self, job: ControlNetQueueEntry) -> List[Image.Image]:
         logger.debug(f"Requested ControlNet: {job.data.controlnet}")
         pipe = self.create_pipe(
             controlnet=job.data.controlnet,
-            scheduler=(job.data.scheduler, job.data.use_karras_sigmas),
+            scheduler=(job.data.scheduler, job.data.sigmas),
             sampler_settings=job.data.sampler_settings,
         )
 

core/inference/utilities/scheduling.py

@@ -11,7 +11,7 @@
 from core.config import config
 from core.inference.utilities.philox import PhiloxGenerator
 from core.scheduling import KdiffusionSchedulerAdapter, create_sampler
-from core.types import PyTorchModelType
+from core.types import PyTorchModelType, SigmaType
 from core.utils import unwrap_enum
 
 logger = logging.getLogger(__name__)
@@ -77,7 +77,7 @@ def change_scheduler(
     model: Optional[PyTorchModelType],
     scheduler: Union[str, KarrasDiffusionSchedulers],
     configuration: Optional[Dict] = None,
-    use_karras_sigmas: bool = False,
+    sigma_type: SigmaType = "",
     sampler_settings: Optional[Dict] = None,
 ) -> SchedulerMixin:
     "Change the scheduler of the model"
@@ -97,9 +97,9 @@ def change_scheduler(
 
         if scheduler.value in [10, 11]:
             logger.debug(
-                f"Loading scheduler {new_scheduler.__class__.__name__} with config karras_sigmas={use_karras_sigmas}"
+                f"Loading scheduler {new_scheduler.__class__.__name__} with config sigmas={sigma_type}"
             )
-            new_scheduler = new_scheduler.from_config(config=configuration, use_karras_sigmas=use_karras_sigmas)  # type: ignore
+            new_scheduler = new_scheduler.from_config(config=configuration, use_karras_sigmas=sigma_type == "")  # type: ignore
         else:
             new_scheduler = new_scheduler.from_config(config=configuration)  # type: ignore
     else:
@@ -109,7 +109,7 @@ def change_scheduler(
             alphas_cumprod=sched.alphas_cumprod,  # type: ignore
             denoiser_enable_quantization=True,
             sampler=scheduler,
-            karras_sigma_scheduler=use_karras_sigmas,
+            sigma_type=sigma_type,
             eta_noise_seed_delta=0,
             sigma_always_discard_next_to_last=False,
             sigma_use_old_karras_scheduler=False,

core/scheduling/scheduling.py

@@ -4,6 +4,8 @@
 import torch
 from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
 
+from core.types import SigmaType
+
 from .adapter.k_adapter import KdiffusionSchedulerAdapter
 from .adapter.unipc_adapter import UnipcSchedulerAdapter
 from .custom.dpmpp_2m import sample_dpmpp_2mV2
@@ -96,7 +98,7 @@ def create_sampler(
     dtype: torch.dtype,
     eta_noise_seed_delta: Optional[float] = None,
     denoiser_enable_quantization: bool = False,
-    karras_sigma_scheduler: bool = False,
+    sigma_type: SigmaType = "",
     sigma_use_old_karras_scheduler: bool = False,
     sigma_always_discard_next_to_last: bool = False,
     sigma_rho: Optional[float] = None,
@@ -135,7 +137,7 @@ def create_sampler(
         )
     else:
         scheduler_name = sampler_tuple[2].get(
-            "scheduler", "karras" if karras_sigma_scheduler else None
+            "scheduler", "karras" if sigma_type == "karras" else None
         )
         if scheduler_name == "karras" and sigma_use_old_karras_scheduler:
             sigma_min = 0.1

core/types.py

@@ -15,6 +15,7 @@
 from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
 
 InferenceBackend = Literal["PyTorch", "AITemplate", "ONNX"]
+SigmaType = Literal["", "karras", "automatic", "exponential", "polyexponential", "vp"]
 Backend = Literal[
     "PyTorch",
     "AITemplate",
@@ -63,7 +64,7 @@ class Txt2imgData:
     steps: int = field(default=25)
     guidance_scale: float = field(default=7)
     self_attention_scale: float = field(default=0.0)
-    use_karras_sigmas: bool = field(default=False)
+    sigmas: SigmaType = field(default="")
     seed: int = field(default=0)
     batch_size: int = field(default=1)
     batch_count: int = field(default=1)
@@ -84,7 +85,7 @@ class Img2imgData:
     steps: int = field(default=25)
     guidance_scale: float = field(default=7)
     self_attention_scale: float = field(default=0.0)
-    use_karras_sigmas: bool = field(default=False)
+    sigmas: SigmaType = field(default="")
     seed: int = field(default=0)
     batch_size: int = field(default=1)
     batch_count: int = field(default=1)
@@ -107,7 +108,7 @@ class InpaintData:
     steps: int = field(default=25)
     guidance_scale: float = field(default=7)
     self_attention_scale: float = field(default=0.0)
-    use_karras_sigmas: bool = field(default=False)
+    sigmas: SigmaType = field(default="")
     seed: int = field(default=0)
     batch_size: int = field(default=1)
     batch_count: int = field(default=1)
@@ -128,7 +129,7 @@ class ControlNetData:
     height: int = field(default=512)
     steps: int = field(default=25)
     guidance_scale: float = field(default=7)
-    use_karras_sigmas: bool = field(default=False)
+    sigmas: SigmaType = field(default="")
     seed: int = field(default=0)
     batch_size: int = field(default=1)
     batch_count: int = field(default=1)

frontend/dist/assets/AboutView.js

@@ -1,4 +1,4 @@
-import { _ as _export_sfc, o as openBlock, a as createElementBlock, i as createBaseVNode } from "./index.js";
+import { _ as _export_sfc, o as openBlock, a as createElementBlock, b as createBaseVNode } from "./index.js";
 const _sfc_main = {};
 const _hoisted_1 = { class: "about" };
 const _hoisted_2 = /* @__PURE__ */ createBaseVNode("h1", null, "This is an about page", -1);

frontend/dist/assets/AccelerateView.js

@@ -1,4 +1,4 @@
-import { T as cB, ae as cM, ad as c, au as cE, aU as iconSwitchTransition, af as cNotM, d as defineComponent, V as useConfig, aj as useRtl, W as useTheme, a6 as provide, s as h, ax as flatten, ay as getSlot, S as createInjectionKey, bh as stepsLight, U as inject, a$ as throwError, c as computed, Z as useThemeClass, aw as resolveWrappedSlot, ak as resolveSlot, aJ as NIconSwitchTransition, ao as createKey, a2 as call, al as NBaseIcon, bi as FinishedIcon, bj as ErrorIcon, x as useMessage, v as useState, r as ref, o as openBlock, a as createElementBlock, b as createVNode, w as withCtx, e as unref, B as NSpace, N as NCard, i as createBaseVNode, p as NSelect, g as NButton, h as createTextVNode, m as NModal, z as serverUrl, u as useSettings, k as createBlock, L as NTabPane, M as NTabs } from "./index.js";
+import { T as cB, ae as cM, ad as c, au as cE, aU as iconSwitchTransition, af as cNotM, d as defineComponent, V as useConfig, aj as useRtl, W as useTheme, a6 as provide, s as h, ax as flatten, ay as getSlot, S as createInjectionKey, bh as stepsLight, U as inject, a$ as throwError, c as computed, Z as useThemeClass, aw as resolveWrappedSlot, ak as resolveSlot, aJ as NIconSwitchTransition, ao as createKey, a2 as call, al as NBaseIcon, bi as FinishedIcon, bj as ErrorIcon, x as useMessage, v as useState, r as ref, o as openBlock, a as createElementBlock, e as createVNode, w as withCtx, f as unref, B as NSpace, N as NCard, b as createBaseVNode, p as NSelect, h as NButton, i as createTextVNode, m as NModal, z as serverUrl, u as useSettings, k as createBlock, L as NTabPane, M as NTabs } from "./index.js";
 import { N as NSlider, a as NSwitch } from "./Switch.js";
 import { N as NInputNumber } from "./InputNumber.js";
 const style = cB("steps", `

frontend/dist/assets/CloudUpload.js

@@ -1,4 +1,4 @@
-import { d as defineComponent, o as openBlock, a as createElementBlock, i as createBaseVNode } from "./index.js";
+import { d as defineComponent, o as openBlock, a as createElementBlock, b as createBaseVNode } from "./index.js";
 const _hoisted_1 = {
   xmlns: "http://www.w3.org/2000/svg",
   "xmlns:xlink": "http://www.w3.org/1999/xlink",

frontend/dist/assets/ExtraView.js

@@ -1,7 +1,7 @@
 import { _ as _sfc_main$2 } from "./GenerateSection.vue_vue_type_script_setup_true_lang.js";
 import { _ as _sfc_main$3 } from "./ImageOutput.vue_vue_type_script_setup_true_lang.js";
 import { I as ImageUpload } from "./ImageUpload.js";
-import { d as defineComponent, v as useState, u as useSettings, x as useMessage, c as computed, o as openBlock, a as createElementBlock, b as createVNode, w as withCtx, e as unref, A as NGi, N as NCard, B as NSpace, i as createBaseVNode, p as NSelect, n as NTooltip, h as createTextVNode, H as NGrid, P as upscalerOptions, z as serverUrl, J as pushScopeId, K as popScopeId, _ as _export_sfc, k as createBlock, L as NTabPane, M as NTabs } from "./index.js";
+import { d as defineComponent, v as useState, u as useSettings, x as useMessage, c as computed, o as openBlock, a as createElementBlock, e as createVNode, w as withCtx, f as unref, A as NGi, N as NCard, B as NSpace, b as createBaseVNode, p as NSelect, n as NTooltip, i as createTextVNode, H as NGrid, P as upscalerOptions, z as serverUrl, J as pushScopeId, K as popScopeId, _ as _export_sfc, k as createBlock, L as NTabPane, M as NTabs } from "./index.js";
 import { N as NSlider } from "./Switch.js";
 import { N as NInputNumber } from "./InputNumber.js";
 import "./SendOutputTo.vue_vue_type_script_setup_true_lang.js";

frontend/dist/assets/GenerateSection.vue_vue_type_script_setup_true_lang.js

@@ -1,4 +1,4 @@
-import { d as defineComponent, o as openBlock, a as createElementBlock, i as createBaseVNode, v as useState, u as useSettings, r as ref, ba as onMounted, y as onUnmounted, z as serverUrl, k as createBlock, w as withCtx, b as createVNode, e as unref, A as NGi, g as NButton, q as NIcon, h as createTextVNode, H as NGrid, bQ as NAlert, G as createCommentVNode, N as NCard } from "./index.js";
+import { d as defineComponent, o as openBlock, a as createElementBlock, b as createBaseVNode, v as useState, u as useSettings, r as ref, ba as onMounted, y as onUnmounted, z as serverUrl, k as createBlock, w as withCtx, e as createVNode, f as unref, A as NGi, h as NButton, q as NIcon, i as createTextVNode, H as NGrid, bQ as NAlert, G as createCommentVNode, N as NCard } from "./index.js";
 const _hoisted_1$1 = {
   xmlns: "http://www.w3.org/2000/svg",
   "xmlns:xlink": "http://www.w3.org/1999/xlink",

frontend/dist/assets/GridOutline.js

@@ -1,4 +1,4 @@
-import { d as defineComponent, o as openBlock, a as createElementBlock, i as createBaseVNode } from "./index.js";
+import { d as defineComponent, o as openBlock, a as createElementBlock, b as createBaseVNode } from "./index.js";
 const _hoisted_1 = {
   xmlns: "http://www.w3.org/2000/svg",
   "xmlns:xlink": "http://www.w3.org/1999/xlink",