AnimateDiff + Faster-Diffusion & DeepCache

.github/CODEOWNERS

@@ -0,0 +1,21 @@
+# Most of the stuff is Staxs
+# core is somewhat general
+* @Stax124
+/core/ @Stax124 @gabe56f
+/frontend/ @Stax124 @gabe56f
+
+# Stax-specific
+/core/* @Stax124
+/core/inference/esrgan/ @Stax124
+
+# Gabe-specific stuff
+/libs/ @gabe56f
+/core/scheduling/ @gabe56f
+/core/optimizations/ @gabe56f
+/core/inference/injectables/ @gabe56f
+/core/inference/utilities/kohya_hires.py @gabe56f
+/core/inference/utilities/anisotropic.py @gabe56f
+/core/inference/utilities/cfg.py @gabe56f
+/core/inference/utilities/sag/ @gabe56f
+/core/inference/utilities/prompt_expansion/ @gabe56f
+/core/inference/onnx/ @gabe56f

.github/workflows/docker_build_tag.yml

@@ -0,0 +1,27 @@
+name: Docker Build on Tag Push
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          file: docker/cuda/dockerfile
+          push: true
+          tags: ${{ secrets.DOCKERHUB_USERNAME }}/volta:${{ github.ref_name }}-cuda

.gitignore

@@ -48,12 +48,12 @@ cover/
 # Diffusers convert files
 out
 traced_unet/
-onnx
+/onnx
 converted
 
 # Docker
-test.docker-compose.yml
-test-no-mount.docker-compose.yml
+*test.docker-compose.yml
+*test-no-mount.docker-compose.yml
 
 # Docs
 node_modules/

.vscode/launch.json

@@ -1,12 +1,12 @@
 {
-	"configurations": [
-		{
-			"name": "Python: File",
-			"type": "python",
-			"request": "launch",
-			"program": "main.py",
-			"args": ["--log-level=DEBUG"],
-			"justMyCode": false
-		}
-	]
+  "configurations": [
+    {
+      "name": "VoltaML API Debug",
+      "type": "python",
+      "request": "launch",
+      "program": "main.py",
+      "args": ["--log-level=DEBUG"],
+      "justMyCode": false
+    }
+  ]
 }

.vscode/settings.json

@@ -4,11 +4,8 @@
 	"python.testing.pytestEnabled": true,
 	"python.analysis.typeCheckingMode": "basic",
 	"python.languageServer": "Pylance",
-	"rust-analyzer.linkedProjects": [
-		"./manager/Cargo.toml"
-	],
+	"rust-analyzer.linkedProjects": ["./manager/Cargo.toml"],
 	"[python]": {
 		"editor.defaultFormatter": "ms-python.black-formatter"
-	},
-	"python.formatting.provider": "none"
+	}
 }

api/app.py

@@ -19,7 +19,9 @@
 from api.websockets.data import Data
 from api.websockets.notification import Notification
 from core import shared
+from core.files import get_full_model_path
 from core.types import InferenceBackend
+from core.utils import determine_model_type
 
 logger = logging.getLogger(__name__)
 
@@ -47,12 +49,33 @@ async def validation_exception_handler(_request: Request, exc: RequestValidation
 
     logger.debug(exc)
 
+    if exc._error_cache is not None and exc._error_cache[0]["loc"][0] == "body":
+        from core.config._config import Configuration
+
+        default_value = Configuration()
+        keys = [str(i) for i in exc._error_cache[0]["loc"][1:]]  # type: ignore
+        current_value = exc._error_cache[0]["ctx"]["given"]  # type: ignore
+
+        # Traverse the config object to find the correct value
+        for key in keys:
+            default_value = getattr(default_value, key)
+
+        websocket_manager.broadcast_sync(
+            data=Data(
+                data={
+                    "default_value": default_value,
+                    "key": keys,
+                    "current_value": current_value,
+                },
+                data_type="incorrect_settings_value",
+            )
+        )
+
     try:
-        why = str(exc).split(":")[1].strip()
-        await websocket_manager.broadcast(
+        websocket_manager.broadcast_sync(
             data=Notification(
                 severity="error",
-                message=f"Validation error: {why}",
+                message="Validation error",
                 title="Validation Error",
             )
         )
@@ -130,7 +153,9 @@ async def startup_event():
         for model in config.api.autoloaded_models:
             if model in [i.path for i in all_models]:
                 backend: InferenceBackend = [i.backend for i in all_models if i.path == model][0]  # type: ignore
-                gpu.load_model(model, backend)
+                model_type = determine_model_type(get_full_model_path(model))[1]
+
+                gpu.load_model(model, backend, type=model_type)
             else:
                 logger.warning(f"Autoloaded model {model} not found, skipping")
 

api/routes/models.py

@@ -20,10 +20,11 @@
     DeleteModelRequest,
     InferenceBackend,
     ModelResponse,
+    PyTorchModelBase,
     TextualInversionLoadRequest,
     VaeLoadRequest,
 )
-from core.utils import download_file
+from core.utils import determine_model_type, download_file
 
 router = APIRouter(tags=["models"])
 logger = logging.getLogger(__name__)
@@ -62,9 +63,11 @@ def list_loaded_models() -> List[ModelResponse]:
 
     loaded_models = []
     for model_id in gpu.loaded_models:
+        name, type_, stage = determine_model_type(get_full_model_path(model_id))
+
         loaded_models.append(
             ModelResponse(
-                name=Path(model_id).name
+                name=name
                 if (".ckpt" in model_id) or (".safetensors" in model_id)
                 else model_id,
                 backend=gpu.loaded_models[model_id].backend,
@@ -75,6 +78,8 @@ def list_loaded_models() -> List[ModelResponse]:
                     "textual_inversions", []
                 ),
                 valid=True,
+                stage=stage,
+                type=type_,
             )
         )
 
@@ -92,11 +97,12 @@ def list_available_models() -> List[ModelResponse]:
 def load_model(
     model: str,
     backend: InferenceBackend,
+    type: PyTorchModelBase,
 ):
     "Loads a model into memory"
 
     try:
-        gpu.load_model(model, backend)
+        gpu.load_model(model, backend, type)
 
         websocket_manager.broadcast_sync(data=Data(data_type="refresh_models", data={}))
     except torch.cuda.OutOfMemoryError:  # type: ignore
@@ -106,7 +112,7 @@ def load_model(
 
 
 @router.post("/unload")
-async def unload_model(model: str):
+def unload_model(model: str):
     "Unloads a model from memory"
 
     gpu.unload(model)
@@ -125,7 +131,7 @@ def unload_all_models():
 
 
 @router.post("/load-vae")
-async def load_vae(req: VaeLoadRequest):
+def load_vae(req: VaeLoadRequest):
     "Load a VAE into a model"
 
     gpu.load_vae(req)
@@ -134,7 +140,7 @@ async def load_vae(req: VaeLoadRequest):
 
 
 @router.post("/load-textual-inversion")
-async def load_textual_inversion(req: TextualInversionLoadRequest):
+def load_textual_inversion(req: TextualInversionLoadRequest):
     "Load a LoRA model into a model"
 
     gpu.load_textual_inversion(req)
@@ -143,15 +149,15 @@ async def load_textual_inversion(req: TextualInversionLoadRequest):
 
 
 @router.post("/memory-cleanup")
-async def cleanup():
+def cleanup():
     "Free up memory manually"
 
     gpu.memory_cleanup()
     return {"message": "Memory cleaned up"}
 
 
 @router.post("/download")
-async def download_model(model: str):
+def download_model(model: str):
     "Download a model to the cache"
 
     gpu.download_huggingface_model(model)
@@ -243,7 +249,7 @@ def delete_model(req: DeleteModelRequest):
 
 @router.post("/download-model")
 def download_checkpoint(
-    link: str, model_type: Literal["Checkpoint", "TextualInversion", "LORA"]
+    link: str, model_type: Literal["Checkpoint", "TextualInversion", "LORA", "VAE"]
 ) -> str:
     "Download a model from a link and return the path to the downloaded file."
 
@@ -254,7 +260,12 @@ def download_checkpoint(
         folder = "textual-inversion"
     elif mtype == "lora":
         folder = "lora"
+    elif mtype == "vae":
+        folder = "vae"
     else:
         raise ValueError(f"Unknown model type {mtype}")
 
-    return download_file(link, Path("data") / folder, True).as_posix()
+    saved_path = download_file(link, Path("data") / folder, True).as_posix()
+    websocket_manager.broadcast_sync(Data(data_type="refresh_models", data={}))
+
+    return saved_path

api/routes/outputs.py

@@ -12,7 +12,7 @@
 thread_pool = ThreadPoolExecutor()
 logger = logging.getLogger(__name__)
 
-valid_extensions = ["png", "jpeg", "webp"]
+valid_extensions = ["png", "jpeg", "webp", "gif"]
 
 
 def sort_images(images: List[Dict[str, Any]]) -> List[Dict[str, Any]]:

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

@@ -1,14 +1,16 @@
 from pathlib import Path
 
-from diffusers.utils.constants import DIFFUSERS_CACHE
+from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE as DIFFUSERS_CACHE
 
-from .config import (
+from ._config import (
     Configuration,
-    Img2ImgConfig,
-    Txt2ImgConfig,
     load_config,
     save_config,
 )
+from .default_settings import (
+    Txt2ImgConfig,
+    Img2ImgConfig,
+)
 
 config = load_config()
 

core/config/_config.py

@@ -0,0 +1,80 @@
+import logging
+from dataclasses import Field, dataclass, field, fields
+
+from dataclasses_json import CatchAll, DataClassJsonMixin, Undefined, dataclass_json
+
+from core.config.samplers.sampler_config import SamplerConfig
+
+from .api_settings import APIConfig
+from .bot_settings import BotConfig
+from .default_settings import (
+    AITemplateConfig,
+    ControlNetConfig,
+    Img2ImgConfig,
+    InpaintingConfig,
+    ONNXConfig,
+    Txt2ImgConfig,
+    UpscaleConfig,
+)
+from .flags_settings import FlagsConfig
+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)
+    sampler_config: SamplerConfig = field(default_factory=SamplerConfig)
+    flags: FlagsConfig = field(default_factory=FlagsConfig)
+    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