Tiled Diffusion / Canvas-based (multi-prompt) generation

.vscode/settings.json

@@ -1,11 +1,14 @@
 {
-	"python.linting.pylintEnabled": true,
-	"python.linting.enabled": true,
-	"python.testing.pytestArgs": ["."],
-	"python.testing.unittestEnabled": false,
-	"python.testing.pytestEnabled": true,
-	"python.analysis.typeCheckingMode": "basic",
-	"python.formatting.provider": "black",
-	"python.languageServer": "Pylance",
-	"rust-analyzer.linkedProjects": ["./manager/Cargo.toml"]
+  "python.linting.pylintEnabled": true,
+  "python.linting.enabled": true,
+  "python.testing.pytestArgs": ["."],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true,
+  "python.analysis.typeCheckingMode": "basic",
+  "python.formatting.provider": "none",
+  "python.languageServer": "Pylance",
+  "rust-analyzer.linkedProjects": ["./manager/Cargo.toml"],
+  "[python]": {
+    "editor.defaultFormatter": "ms-python.black-formatter"
+  }
 }

api/routes/settings.py

@@ -4,7 +4,7 @@
 from fastapi import APIRouter
 
 from core import config
-from core.config.config import update_config
+from core.config._config import update_config
 
 router = APIRouter(tags=["settings"])
 

core/config/__init__.py

@@ -2,7 +2,7 @@
 
 from diffusers.utils.constants import DIFFUSERS_CACHE
 
-from .config import (
+from ._config import (
     Configuration,
     Img2ImgConfig,
     Txt2ImgConfig,

core/config/_config.py

@@ -0,0 +1,75 @@
+import logging
+from dataclasses import Field, dataclass, field, fields
+
+from dataclasses_json import CatchAll, DataClassJsonMixin, Undefined, dataclass_json
+
+from .api_settings import APIConfig
+from .bot_settings import BotConfig
+from .default_settings import (
+    Txt2ImgConfig,
+    Img2ImgConfig,
+    InpaintingConfig,
+    ControlNetConfig,
+    UpscaleConfig,
+    AITemplateConfig,
+    ONNXConfig,
+)
+from .frontend_settings import FrontendConfig
+from .interrogator_settings import InterrogatorConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass_json(undefined=Undefined.INCLUDE)
+@dataclass
+class Configuration(DataClassJsonMixin):
+    "Main configuration class for the application"
+
+    txt2img: Txt2ImgConfig = field(default_factory=Txt2ImgConfig)
+    img2img: Img2ImgConfig = field(default_factory=Img2ImgConfig)
+    inpainting: InpaintingConfig = field(default_factory=InpaintingConfig)
+    controlnet: ControlNetConfig = field(default_factory=ControlNetConfig)
+    upscale: UpscaleConfig = field(default_factory=UpscaleConfig)
+    api: APIConfig = field(default_factory=APIConfig)
+    interrogator: InterrogatorConfig = field(default_factory=InterrogatorConfig)
+    aitemplate: AITemplateConfig = field(default_factory=AITemplateConfig)
+    onnx: ONNXConfig = field(default_factory=ONNXConfig)
+    bot: BotConfig = field(default_factory=BotConfig)
+    frontend: FrontendConfig = field(default_factory=FrontendConfig)
+    extra: CatchAll = field(default_factory=dict)
+
+
+def save_config(config: Configuration):
+    "Save the configuration to a file"
+
+    logger.info("Saving configuration to data/settings.json")
+
+    with open("data/settings.json", "w", encoding="utf-8") as f:
+        f.write(config.to_json(ensure_ascii=False, indent=4))
+
+
+def update_config(config: Configuration, new_config: Configuration):
+    "Update the configuration with new values instead of overwriting the pointer"
+
+    for cls_field in fields(new_config):
+        assert isinstance(cls_field, Field)
+        setattr(config, cls_field.name, getattr(new_config, cls_field.name))
+
+
+def load_config():
+    "Load the configuration from a file"
+
+    logger.info("Loading configuration from data/settings.json")
+
+    try:
+        with open("data/settings.json", "r", encoding="utf-8") as f:
+            config = Configuration.from_json(f.read())
+            logger.info("Configuration loaded from data/settings.json")
+            return config
+
+    except FileNotFoundError:
+        logger.info("data/settings.json not found, creating a new one")
+        config = Configuration()
+        save_config(config)
+        logger.info("Configuration saved to data/settings.json")
+        return config

core/config/api_settings.py

@@ -0,0 +1,110 @@
+from dataclasses import dataclass, field
+from typing import List, Literal, Union
+
+import torch
+
+
+@dataclass
+class APIConfig:
+    "Configuration for the API"
+
+    # Autoload
+    autoloaded_textual_inversions: List[str] = field(default_factory=list)
+
+    # Websockets and intervals
+    websocket_sync_interval: float = 0.02
+    websocket_perf_interval: float = 1.0
+
+    # TomeSD
+    use_tomesd: bool = False  # really extreme, probably will have to wait around until tome improves a bit
+    tomesd_ratio: float = 0.25  # had to tone this down, 0.4 is too big of a context loss even on short prompts
+    tomesd_downsample_layers: Literal[1, 2, 4, 8] = 1
+
+    image_preview_delay: float = 2.0
+
+    # General optimizations
+    autocast: bool = False
+    attention_processor: Literal[
+        "xformers", "sdpa", "cross-attention", "subquadratic", "multihead"
+    ] = "sdpa"
+    subquadratic_size: int = 512
+    attention_slicing: Union[int, Literal["auto", "disabled"]] = "disabled"
+    channels_last: bool = True
+    vae_slicing: bool = True
+    vae_tiling: bool = False
+    trace_model: bool = False
+    clear_memory_policy: Literal["always", "after_disconnect", "never"] = "always"
+    offload: bool = False
+    data_type: Literal["float32", "float16", "bfloat16"] = "float16"
+
+    # CUDA specific optimizations
+    reduced_precision: bool = False
+    cudnn_benchmark: bool = False
+    deterministic_generation: bool = False
+
+    # Device settings
+    device_id: int = 0
+    device_type: Literal["cpu", "cuda", "mps", "directml", "intel", "vulkan"] = "cuda"
+
+    # Critical
+    enable_shutdown: bool = True
+
+    # VAE
+    upcast_vae: bool = False
+
+    # CLIP
+    clip_skip: int = 1
+    clip_quantization: Literal["full", "int8", "int4"] = "full"
+
+    huggingface_style_parsing: bool = False
+
+    # Saving
+    save_path_template: str = "{folder}/{prompt}/{id}-{index}.{extension}"
+    image_extension: Literal["png", "webp", "jpeg"] = "png"
+    image_quality: int = 95
+    image_return_format: Literal["bytes", "base64"] = "base64"
+
+    # Grid
+    disable_grid: bool = False
+
+    # Torch compile
+    torch_compile: bool = False
+    torch_compile_fullgraph: bool = False
+    torch_compile_dynamic: bool = False
+    torch_compile_backend: str = "inductor"
+    torch_compile_mode: Literal[
+        "default",
+        "reduce-overhead",
+        "max-autotune",
+    ] = "reduce-overhead"
+
+    @property
+    def dtype(self):
+        "Return selected data type"
+        if self.data_type == "bfloat16":
+            return torch.bfloat16
+        if self.data_type == "float16":
+            return torch.float16
+        return torch.float32
+
+    @property
+    def device(self):
+        "Return the device"
+
+        if self.device_type == "intel":
+            from core.inference.functions import is_ipex_available
+
+            return torch.device("xpu" if is_ipex_available() else "cpu")
+
+        if self.device_type in ["cpu", "mps"]:
+            return torch.device(self.device_type)
+
+        if self.device_type in ["vulkan", "cuda"]:
+            return torch.device(f"{self.device_type}:{self.device_id}")
+
+        if self.device_type == "directml":
+            import torch_directml  # pylint: disable=import-error
+
+            return torch_directml.device()
+        else:
+            raise ValueError(f"Device type {self.device_type} not supported")

core/config/bot_settings.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
+
+
+@dataclass
+class BotConfig:
+    "Configuration for the bot"
+
+    default_scheduler: KarrasDiffusionSchedulers = (
+        KarrasDiffusionSchedulers.DPMSolverSinglestepScheduler
+    )
+    verbose: bool = False
+    use_default_negative_prompt: bool = True