v0.2 (#88)

Merge features from the experimental branch into the main branch
(stable)

---------

Co-authored-by: Márton Kissik <[email protected]>
Co-authored-by: aaronsantiago <[email protected]>
Co-authored-by: MiningP <[email protected]>
Co-authored-by: Samuel Corsi-House <[email protected]>

Author: 5 people

.github/workflows/yarn_build.yml renamed to .github/workflows/frontend_build.yml

@@ -6,6 +6,9 @@ name: Frontend Build Test
 on:
   push:
     branches: [main, experimental]
+    paths:
+      - "frontend/**"
+      - ".github/workflows/*.yml"
   pull_request:
     branches: [main, experimental]
 
@@ -24,10 +27,10 @@ jobs:
       - name: Run install
         uses: borales/actions-yarn@v4
         with:
-          cmd: install # will run `yarn install` command
-          dir: "frontend" # will run `yarn install` in `frontend` sub folder
+          cmd: install
+          dir: "frontend"
       - name: Build production bundle
         uses: borales/actions-yarn@v4
         with:
-          cmd: build # will run `yarn build:prod` command
-          dir: "frontend" # will run `yarn build` in `frontend` sub folder
+          cmd: build
+          dir: "frontend"

.github/workflows/manager_build.yml

@@ -5,25 +5,37 @@ on:
     branches: ["main", "experimental"]
     paths:
       - "**.rs"
+      - ".github/workflows/**"
   pull_request:
     branches: ["main", "experimental"]
     paths:
       - "**.rs"
+      - ".github/workflows/**"
+  workflow_dispatch:
+    inputs:
+      branch:
+        description: "Branch to build"
+        required: true
+        default: "main"
 
 env:
   CARGO_TERM_COLOR: always
 
 jobs:
   build:
-    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+
+    runs-on: ${{ matrix.os }}
 
     steps:
       - uses: actions/checkout@v3
       - name: Build
-        run: cd manager && cargo build --verbose --release
+        run: cd manager && cargo build --release
       - name: Output binary
         uses: actions/upload-artifact@v3
         with:
           name: volta-manager
-          path: manager/target/release/voltaml-manager
+          path: manager/target/release/voltaml-manager*
           retention-days: 3

.gitignore

@@ -74,7 +74,6 @@ docs/.vitepress/dist
 /engine
 /static/output
 /outputs
-/frontend/yarn.lock
 /testing
 /typings
 external
@@ -88,3 +87,11 @@ core/submodules
 
 # Manager
 manager/target
+voltaml-manager
+voltaml-manager.exe
+
+# Profiler stuff
+profile.html
+
+# dotenv file
+.env

README.md

@@ -53,7 +53,7 @@
 </div>
 
 <hr>
-<h3 align="center">Made with ❤️ by <a href="https://github.com/Stax124/">Stax124</a></h3>
+<h3 align="center">Made with ❤️ by <a href="https://github.com/Stax124/">Stax124</a> and the community</h3>
 <hr>
 
 <br />
@@ -67,7 +67,6 @@
   - [Speed comparison](#speed-comparison)
   - [Installation](#installation)
 - [Contributing](#contributing)
-  - [Code of Conduct](#code-of-conduct)
 - [License](#license)
 - [Contact](#contact)
 
@@ -136,17 +135,11 @@
 
 ## Speed comparison
 
-The below benchmarks have been done for generating a 512x512 image, batch size 1 for 50 iterations.
-
-| GPU (it/s) | T4  | A10  | A100 | 4090 | 4080 | 3090 | 2080Ti | 3050 |
-| ---------- | --- | ---- | ---- | ---- | ---- | ---- | ------ | ---- |
-| PyTorch    | 4.3 | 8.8  | 15.1 | 19   | 15.5 | 11   | 8      | 4.1  |
-| xFormers   | 5.5 | 15.6 | 27.5 | 28   | 20.2 | 15.7 | N/A    | 5.1  |
-| AITemplate | N/A | 23   | N/A  | N/A  | 40.5 | N/A  | N/A    | 10.2 |
+Please refer to this [table](https://voltaml.github.io/voltaML-fast-stable-diffusion/getting-started/introduction#speed-comparison). Data had a small sample size and was usually collected on a single machine. Your results may vary.
 
 ## Installation
 
-Please see the [documentation](https://voltaml.github.io/voltaML-fast-stable-diffusion/installation/docker.html) for installation instructions.
+Please see the [documentation](https://voltaml.github.io/voltaML-fast-stable-diffusion/installation/local) for installation instructions.
 
 # Contributing
 
@@ -158,10 +151,6 @@ Contributions are always welcome!
 
 See `contributing.md` for ways to get started.
 
-## Code of Conduct
-
-Please read the [Code of Conduct](https://github.com/VoltaML/voltaML-fast-stable-diffusion/blob/master/CODE_OF_CONDUCT.md)
-
 # License
 
 Distributed under the <b>GPL v3</b>. See [License](https://github.com/VoltaML/voltaML-fast-stable-diffusion/blob/experimental/License) for more information.

api/app.py

@@ -89,11 +89,31 @@ async def custom_http_exception_handler(_request, _exc):
 async def startup_event():
     "Prepare the event loop for other asynchronous tasks"
 
+    # Inject the logger
+    from rich.logging import RichHandler
+
+    # Disable duplicate logger
+    logging.getLogger("uvicorn").handlers = []
+
+    for logger_ in ("uvicorn.access", "uvicorn.error", "fastapi"):
+        l = logging.getLogger(logger_)
+        handler = RichHandler(rich_tracebacks=True, show_time=False)
+        handler.setFormatter(
+            logging.Formatter(
+                fmt="%(asctime)s | %(name)s » %(message)s", datefmt="%H:%M:%S"
+            )
+        )
+        l.handlers = [handler]
+
     shared.asyncio_loop = asyncio.get_event_loop()
 
-    asyncio.create_task(websocket_manager.sync_loop())
+    sync_task = asyncio.create_task(websocket_manager.sync_loop())
     logger.info("Started WebSocketManager sync loop")
-    asyncio.create_task(websocket_manager.perf_loop())
+    perf_task = asyncio.create_task(websocket_manager.perf_loop())
+
+    shared.asyncio_tasks.append(sync_task)
+    shared.asyncio_tasks.append(perf_task)
+
     logger.info("Started WebSocketManager performance monitoring loop")
     logger.info("UI Available at: http://localhost:5003/")
 
@@ -106,18 +126,6 @@ async def shutdown_event():
     await websocket_manager.close_all()
 
 
-# Origins that are allowed to access the API
-origins = ["*"]
-
-# Allow CORS for specified origins
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 # Enable FastAPI Analytics if key is provided
 key = os.getenv("FASTAPI_ANALYTICS_KEY")
 if key:
@@ -149,7 +157,7 @@ async def shutdown_event():
 static_app = FastAPI()
 static_app.add_middleware(
     CORSMiddleware,
-    allow_origins=origins,
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -160,3 +168,12 @@ async def shutdown_event():
 static_app.mount("/", StaticFiles(directory="frontend/dist/assets"), name="assets")
 
 app.mount("/assets", static_app)
+
+# Allow CORS for specified origins
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)

api/routes/general.py

@@ -1,8 +1,14 @@
+import logging
+import sys
+
 from fastapi import APIRouter
 
+from api import websocket_manager
+from api.websockets.notification import Notification
 from core import shared
 
 router = APIRouter(tags=["general"])
+logger = logging.getLogger(__name__)
 
 
 @router.post("/interrupt")
@@ -11,3 +17,71 @@ async def interrupt():
 
     shared.interrupt = True
     return {"message": "Interupted"}
+
+
+@router.post("/shutdown")
+async def shutdown():
+    "Shutdown the server"
+
+    from core.config import config
+    from core.shared import uvicorn_loop, uvicorn_server
+
+    if config.api.enable_shutdown:
+        if uvicorn_server is not None:
+            await websocket_manager.broadcast(
+                data=Notification(
+                    message="Shutting down the server",
+                    severity="warning",
+                    title="Shutdown",
+                )
+            )
+            for task in shared.asyncio_tasks:
+                task.cancel()
+            uvicorn_server.force_exit = True
+            logger.debug("Setting force_exit to True")
+
+        assert uvicorn_server is not None
+        await uvicorn_server.shutdown()
+        logger.debug("Unicorn server shutdown")
+
+        assert uvicorn_loop is not None
+        uvicorn_loop.stop()
+        logger.debug("Unicorn loop stopped")
+
+        sys.exit(0)
+
+    else:
+        await websocket_manager.broadcast(
+            data=Notification(
+                message="Shutdown is disabled", severity="error", title="Shutdown"
+            )
+        )
+        return {"message": "Shutdown is disabled"}
+
+
+@router.get("/queue-status")
+async def queue_status():
+    "Get the status of the queue"
+
+    from core.shared_dependent import gpu
+
+    queue = gpu.queue
+
+    return {
+        "jobs": queue.jobs,
+        "concurrent_jobs": queue.concurrent_jobs,
+        "locked": queue.lock.locked(),
+    }
+
+
+@router.post("/queue-clear")
+async def queue_clear():
+    "Clear the queue"
+
+    from core.shared_dependent import gpu
+
+    queue = gpu.queue
+
+    queue.clear()
+
+    return {"message": "Queue cleared"}

api/routes/generate.py

@@ -8,20 +8,21 @@
 from core.shared_dependent import gpu
 from core.types import (
     AITemplateBuildRequest,
+    AITemplateDynamicBuildRequest,
     ControlNetQueueEntry,
     ConvertModelRequest,
     Img2ImgQueueEntry,
     InpaintQueueEntry,
     InterrogatorQueueEntry,
     ONNXBuildRequest,
-    RealESRGANQueueEntry,
     SDUpscaleQueueEntry,
     TRTBuildRequest,
     Txt2ImgQueueEntry,
+    UpscaleQueueEntry,
 )
 from core.utils import convert_bytes_to_image_stream, convert_image_to_base64
 
-router = APIRouter(tags=["txt2img"])
+router = APIRouter(tags=["generate"])
 logger = logging.getLogger(__name__)
 
 
@@ -195,38 +196,27 @@ async def sd_upscale_job(job: SDUpscaleQueueEntry):
         }
 
 
-@router.post("/realesrgan-upscale")
-async def realesrgan_upscale_job(job: RealESRGANQueueEntry):
+@router.post("/upscale")
+async def realesrgan_upscale_job(job: UpscaleQueueEntry):
     "Upscale image with RealESRGAN model"
 
     image_bytes = job.data.image
     assert isinstance(image_bytes, bytes)
     job.data.image = convert_bytes_to_image_stream(image_bytes)
 
     try:
-        images: Union[List[Image.Image], List[str]]
+        image: Image.Image
         time: float
-        images, time = await gpu.generate(job)
+        image, time = await gpu.upscale(job)
     except ModelNotLoadedError:
         raise HTTPException(  # pylint: disable=raise-missing-from
             status_code=400, detail="Model is not loaded"
         )
 
-    if len(images) == 0:
-        return {
-            "time": time,
-            "images": [],
-        }
-    elif isinstance(images[0], str):
-        return {
-            "time": time,
-            "images": images,
-        }
-    else:
-        return {
-            "time": time,
-            "images": [convert_image_to_base64(i) for i in images],  # type: ignore
-        }
+    return {
+        "time": time,
+        "images": convert_image_to_base64(image),  # type: ignore
+    }
 
 
 @router.post("/generate-engine")
@@ -247,6 +237,15 @@ async def generate_aitemplate(request: AITemplateBuildRequest):
     return {"message": "Success"}
 
 
+@router.post("/generate-dynamic-aitemplate")
+async def generate_dynamic_aitemplate(request: AITemplateDynamicBuildRequest):
+    "Generate a TensorRT engine from a local model"
+
+    await gpu.build_dynamic_aitemplate_engine(request)
+
+    return {"message": "Success"}
+
+
 @router.post("/generate-onnx")
 async def generate_onnx(request: ONNXBuildRequest):
     "Generate a TensorRT engine from a local model"