Implement task cancellation

.env.test

@@ -0,0 +1,9 @@
+SECRET_KEY=30a1590ae015eac4908095904d6dc0e086a61d1a6acd6058f0560adfe5b253e5
+PRINT_LINK_WITH_TOKEN=
+BASE_URL=http://localhost:8000
+ENABLE_PUBLIC_ACCESS=1
+ENABLE_SIGNUP=
+REDIS_HOST=redis
+REDIS_PORT=6379
+REDIS_PASSWORD=jdfhefo1e1928h1o2389r123fdjq92edh12pd12ed2d12d
+HUGGINGFACE_TOKEN=

.github/workflows/ci.yaml

@@ -109,4 +109,4 @@ jobs:
       - uses: actions/checkout@v2
 
       - name: E2E docker-compose tests
-        run: make tests-up
+        run: make test

Makefile

@@ -1,8 +1,5 @@
-tests-up:
-	docker-compose up --abort-on-container-exit --exit-code-from e2e_tests --force-recreate --build api redis_worker e2e_tests
-
-tests-run:
-	docker-compose run e2e_tests
+test:
+	docker-compose --env-file .env.test up --abort-on-container-exit --exit-code-from e2e_tests --force-recreate --build api redis_worker e2e_tests
 
 run:
 	docker-compose up --abort-on-container-exit --force-recreate --build api redis_worker

README.md

@@ -1,7 +1,7 @@
 # Stable Diffusion API
 
-[![OpenApi](https://img.shields.io/badge/OpenApi-3.0.2-orange)](https://editor.swagger.io/?url=https://raw.githubusercontent.com/irgolic/stable-diffusion-server/master/openapi.yml)
-[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/irgolic/stable-diffusion-server/blob/master/colab_runner.ipynb)
+[![OpenApi](https://img.shields.io/badge/OpenApi-3.0.2-orange)](https://editor.swagger.io/?url=https://raw.githubusercontent.com/irgolic/stable-diffusion-api/master/openapi.yml)
+[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/irgolic/stable-diffusion-api/blob/master/colab_runner.ipynb)
 [![Discord](https://discordapp.com/api/guilds/1027703335224098857/widget.png?style=shield)](https://discord.gg/UXQfCRpYSC)
 
 Simple backend to serve Txt2Img, Img2Img and Inpainting with any model published on [Hugging Face](https://huggingface.co/models).
@@ -36,7 +36,7 @@ Inpainting also supports:
 ## Usage
 
 Generate any client library from the [OpenApi](
-https://editor.swagger.io/?url=https://raw.githubusercontent.com/irgolic/stable-diffusion-server/master/openapi.yml) specification.
+https://editor.swagger.io/?url=https://raw.githubusercontent.com/irgolic/stable-diffusion-api/master/openapi.yml) specification.
 
 ### Authentication
 
@@ -52,14 +52,17 @@ Registered users can generate their own tokens at `POST /token/{username}`.
 
 ### Synchronous Interface
 
-For convenience, the server provides synchronous endpoints at `GET /txt2img`, `GET /img2img`, and `GET /inpaint`.
+For convenience, the API provides synchronous endpoints at `GET /txt2img`, `GET /img2img`, and `GET /inpaint`.
 
 To print a browser-accessible URL upon startup (i.e., `http://localhost:8000/txt2img?prompt=corgi&steps=5?token=...`), 
 set environment variable `PRINT_LINK_WITH_TOKEN=1` (set by default in `.env.example`).
 
-The server will wait for the model to download and generate an image before returning the request.
+The API will wait for the model to download and generate an image before returning the request.
 It is preferable to use the asynchronous endpoints for production use.
 
+If the connection is dropped (you navigate away from the page),
+the API will automatically cancel the request and free up resources.
+
 ### Asynchronous Interface (recommended)
 
 #### Invocation
@@ -77,11 +80,13 @@ Event types:
 - FinishedEvent (with `blob_url` and `parameters_used`)
 - AbortedEvent (with `reason`)
 
+To cancel a task, `DELETE /task/{task_id}`.
+
 #### Results
 
 The FinishedEvent contains a URL to the generated image.
 Currently only local blob serving is supported;
-make sure to download the image before shutting down the server.
+make sure to download the image before shutting down the API.
 
 ## Installation
 
@@ -96,8 +101,8 @@ Install python 3.10 with your preferred environment creator.
 Install [MiniConda](https://docs.conda.io/en/latest/miniconda.html) and create a new environment with python 3.10.
 
 ```bash
-conda create -n sds python=3.10
-conda activate sds
+conda create -n sda python=3.10
+conda activate sda
 ```
 
 ### Poetry
@@ -138,19 +143,19 @@ make run
 Or invoke processes on multiple machines, starting the API with:
 
 ```bash
-poetry run uvicorn stable_diffusion_server.api.redis_app:app
+poetry run uvicorn stable_diffusion_api.api.redis_app:app
 ```
 
 And the worker(s) with:
 
 ```bash
-poetry run python -m stable_diffusion_server.engine.worker.redis_worker
+poetry run python -m stable_diffusion_api.engine.worker.redis_worker
 ```
 
 ### Single Process
 
 Run the API and worker in a single process. API requests will block until the worker is finished.
 
 ```bash
-poetry run uvicorn stable_diffusion_server.api.in_memory_app:app
+poetry run uvicorn stable_diffusion_api.api.in_memory_app:app
 ```

colab_runner.ipynb

@@ -46,8 +46,8 @@
    "source": [
     "#@markdown # Clone git repo\n",
     "\n",
-    "!git clone https://github.com/irgolic/stable-diffusion-server\n",
-    "%cd /content/stable-diffusion-server"
+    "!git clone https://github.com/irgolic/stable-diffusion-api\n",
+    "%cd /content/stable-diffusion-api"
    ],
    "metadata": {
     "id": "vjZsZgYMQxew",
@@ -140,7 +140,7 @@
    "source": [
     "#@markdown # Install dependencies\n",
     "\n",
-    "%cd /content/stable-diffusion-server\n",
+    "%cd /content/stable-diffusion-api\n",
     "!poetry install\n",
     "\n",
     "# For URL forwarding in Colab\n",
@@ -201,7 +201,7 @@
    "source": [
     "#@markdown # Run Server\n",
     "#@markdown The first time you call /txt2img in your browser it might time out; let the model download and cache, then try again.\n",
-    "%cd /content/stable-diffusion-server\n",
+    "%cd /content/stable-diffusion-api\n",
     "\n",
     "import time\n",
     "\n",
@@ -224,7 +224,7 @@
     "os.environ['BASE_URL'] = local_tunnel_url\n",
     "\n",
     "# run in memory (TODO run multiple processes intermediated by redis)\n",
-    "!poetry run uvicorn stable_diffusion_server.api.in_memory_app:app"
+    "!poetry run uvicorn stable_diffusion_api.api.in_memory_app:app"
    ],
    "metadata": {
     "id": "dSlL_wCSgHJ3",

deployment/run_redis_app.sh

@@ -1,2 +1,2 @@
 #!/bin/bash
-poetry run uvicorn stable_diffusion_server.api.redis_app:app --host 0.0.0.0
+poetry run uvicorn stable_diffusion_api.api.redis_app:app --host 0.0.0.0

deployment/run_redis_worker.sh

@@ -1,2 +1,2 @@
 #!/bin/bash
-poetry run python -m stable_diffusion_server.engine.workers.redis_worker
+poetry run python -m stable_diffusion_api.engine.workers.redis_worker

docker-compose.yaml

@@ -14,7 +14,7 @@ services:
       - ./:/app
   api:
     build: .
-    image: irgolic/stable-diffusion-server
+    image: irgolic/stable-diffusion-api:latest
     depends_on:
       - redis
 #    expose:
@@ -35,7 +35,7 @@ services:
       - ./:/app
     command: deployment/run_redis_app.sh
   redis_worker:
-    image: irgolic/stable-diffusion-server:latest
+    image: irgolic/stable-diffusion-api:latest
     deploy:
       replicas: 5
     depends_on:
@@ -53,7 +53,7 @@ services:
       - ~/.cache/huggingface:/root/.cache/huggingface
     entrypoint: deployment/run_redis_worker.sh
   e2e_tests:
-    image: irgolic/stable-diffusion-server:latest
+    image: irgolic/stable-diffusion-api:latest
     depends_on:
       - redis
       - api

openapi.yml

@@ -788,6 +788,34 @@ paths:
       - OAuth2PasswordBearer: []
       summary: Create Task
   /task/{task_id}:
+    delete:
+      operationId: delete_task_task__task_id__delete
+      parameters:
+      - in: path
+        name: task_id
+        required: true
+        schema:
+          title: Task Id
+          type: string
+      responses:
+        '200':
+          content:
+            application/json:
+              schema: {}
+          description: Successful Response
+        '204':
+          description: Task cancelled
+        '404':
+          description: Task not found
+        '422':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+          description: Validation Error
+      security:
+      - OAuth2PasswordBearer: []
+      summary: Delete Task
     get:
       operationId: poll_task_status_task__task_id__get
       parameters:

pyproject.toml

@@ -1,11 +1,11 @@
 [tool.poetry]
-name = "stable-diffusion-server"
+name = "stable-diffusion-api"
 version = "0.1.0"
 description = ""
 authors = ["Rafael Irgolic <[email protected]>"]
 license = "GNU Affero General Public License"
 readme = "README.md"
-packages = [{include = "stable_diffusion_server"}]
+packages = [{include = "stable_diffusion_api"}]
 
 [tool.poetry.dependencies]
 python = "^3.10"