Skip to content

Commit

Permalink
fixes server and client
Browse files Browse the repository at this point in the history 
adds new instructions to readme
  • Loading branch information
@w4ffl35
w4ffl35 committed Aug 13, 2024
1 parent 34bae6c commit 8f2b1da
Show file tree
Hide file tree
Showing 15 changed files with 441 additions and 426 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,4 @@ tmp/*
build.log
venv
*.egg-info
settings.py
292 changes: 13 additions & 279 deletions README.md
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# RunAI

Run AI allows you to run a threaded Stable Diffusion low level python socket
server.
Run AI allows you to run a LLMs using a socket server.

---

Expand All @@ -11,8 +10,6 @@ server.
- **Sockets**: handles byte packets of an arbitrary size
- **Threaded**: asynchronously handle requests and responses
- **Queue**: requests and responses are handed off to a queue
- **Auto-shutdown**: server automatically shuts down after client disconnects
- Does not save images or logs to disc

---

Expand All @@ -27,291 +24,28 @@ Stable Diffusion (and other AI models) locally. It was designed for use with the
Krita Stable Diffusion plugin, but can work with any interface provided someone
writes a client for it.

### Only uses float16 (half floats)
### Only works with Mistral

If someone wants to build in functionality for float32 I will merge the code but
currently this is not a priority feature.
This library was designed to work with the Mistral model, but it can be expanded
to work with any LLM.

---

## Installation

First run `sh bin/install.sh` to install the required models. These will be
placed in `~/stablediffusion`. See [Stable Diffusion directory structure](#stable-diffusion-directory-structure) for more information.

## Docker

Easiest method

1. [Install docker](https://docs.docker.com/engine/install/)
2. [Install nvdia-container-runtime](https://nvidia.github.io/nvidia-container-runtime/)
3. `sudo apt install nvidia-container-toolkit`
4. Copy `daemon.json` to `/etc/docker/daemon.json` (if you already have a daemon.js file in that directory, just copy the contents)
5. `docker-compose up`

----

### Docker commands

All of the following commands are contained in `/bin/dc`, you can add it to your path or run it directly.

- **Run** `./bin/dc start` run the server
- **Shell** `./bin/dc bash` enter shell
- **Update** `./bin/dc updatereqs` update pip
- **Build** `./bin/dc build` build the server
- **Clean** `./bin/dc /app/bin/clean.sh`

---

### More commands

- Build and start the services `docker-compose up`
- Stop and remove all services `docker-compose down`
- Rebuild all services `docker-compose build`
- List all running containers `docker-compose ps`
- View the output from containers `docker-compose logs`
- Execute a command in a running container `docker-compose exec <service> <command>`
- Replace <service> with the name of the service defined in the docker-compose.yml file, and <command> with the command you want to run.

---

## Bare metal

1. [Install CUDA Toolkit 11.7](https://developer.nvidia.com/cuda-11-7-0-download-archive?target_os=Linux&target_arch=x86_64)
2. [Install miniconda](https://docs.conda.io/en/latest/miniconda.html)
3. Activate environment ``
4. conda activate runai
5. Install requirements `pip install -r requirements.txt`

Create a lib folder in the root of the project:

`mkdir -r lib/torch`

Copy the following into `lib/torch/`:

- `lib/torch/bin/torch_shm_manager`
- `lib/torch/lib/libtorch_global_deps.so`

Your directory structure may differ, but it will likely look something like this:

```
/home/<user>/miniconda3/envs/ksd-build/lib/python3.10/site-packages/torch/bin/torch_shm_manager
/home/<user>/miniconda3/envs/ksd-build/lib/python3.10/site-packages/torch/lib/libtorch_global_deps.so
```

![img.png](img.png)

- git
- conda
- a cuda capable GPU

Build the server
```
./bin/buildlinux.sh
```
The standalone server will be in the `dist` directory

Run the server
```
conda activate runai
python server.py
```

---

## Request structure

Clients establish a connection with the server over a socket and send a JSON
object encoded as a byte string split into packets. An EOM (end of message)
signal is sent to indicate the end of the message.


![img_1.png](img_1.png)

The server assembles the packets, decodes the JSON object and processes the
request. Once processing is complete the server will send a response
back to the client.

It is up to the client to reassemble the packets, decode the byte string to JSON
and handle the message.

---

## Client

For an example client, take a look at the [connect.py file](https://github.com/w4ffl35/krita_stable_diffusion/blob/master/krita_stable_diffusion/connect.py) in the Krita Stable Diffusion [Plugin](https://github.com/w4ffl35/krita_stable_diffusion) which uses this server.

---

## Stable Diffusion directory structure


This is the recommended and default setup for runai

### Linux

Default directory structure for runai Stable Diffusion

#### Base models

These models are required to run Stable Diffusion

- **CLIP** files for CLIP
- **CompVis** safety checker model (used for NSWF filtering)
- **openai** clip-vit-large-patch14 model

```
├── ~/stablediffusuion
   ├── CLIP
   ├── CompVis
   │   ├── stable-diffusion-safety-checker
   ├── openai
      ├── clip-vit-large-patch14
```

#### Diffusers models

These are the base models to run a particular version of Stable Diffusion.

- **runwayml**: Base models for Stable Diffusion v1
- **stabilityai**: Base models for Stable Diffusion v2

```
├── ~/stablediffusuion
   ├── runwayml
      ├── stable-diffusion-inpainting
      ├── stable-diffusion-v1-5
   ├── stabilityai
      ├── stable-diffusion-2-1-base
      ├── stable-diffusion-2-inpainting
```

#### Custom models

- **v1** should be a directory containing models using stable diffusion v1
- **v2** should be a directory containing models using stable diffusion v2

You may place diffusers folders, ckpt and safetensor files in these directories.

```
├── ~/stablediffusuion
   ├── v1
   │   ├── <folder> (diffusers directory)
   │   ├── <file>.ckpt
   │   ├── <file>.safetensor
   ├── v2
      ├── <folder> (diffusers directory)
      ├── <file>.ckpt
      ├── <file>.safetensor
```

### Automatic1111 existing files

If you are using **Automatic1111** you can place your checkpoints in the
webui models folder as you typically would, however the directory structure
which includes v1 models separated from v2 models is required for now.

This allows you to use the same checkpoints for both **Automatic1111 webui**
and this server.

For example, if your `webui` directory looks like this

```
├── /home/USER/stable-diffusion-webui/models/Stable-diffusion
   ├── <some_checkpoint_file>.ckpt
   ├── <some_other_checkpoint_file>.ckpt
   ├── <some_other_checkpoint_file_v2>.ckpt
```

You would reorganize it like this:

```
├── /home/USER/stable-diffusion-webui/models/Stable-diffusion
   ├── v1
      ├── <some_checkpoint_file>.ckpt
      ├── <some_other_checkpoint_file>.ckpt
   ├── v2
      ├── <some_other_checkpoint_file_v2>.ckpt
```

You would then set BASE_DIR to `/home/USER/stable-diffusion-webui/models/Stable-diffusion`

---

### Build

First install `pyinstaller`

`pip install pyinstaller`

Then build the executable

```bash
pip install runai
cp src/runai/default.settings.py src/runai/settings.py
```
./bin/buildlinux.sh
```

Test

```
cd ./dist/runai
./runai
```

This should start a server.

[Connect a client to see if it is working properly](https://github.com/w4ffl35/krita_stable_diffusion)

---

## Running the server

`python server.py`

The following flags and options are available

- `--port` (int) - port to run server on
- `--host` (str) - host to run server on
- `--timeout` - whether to timeout after failing to receive a client connection, pass this flag for true, otherwise the server will not timeout.
- `--packet-size` (int) - size of byte packets to transmit to and from the client
- `--model-base-path` (str) - base directory for checkpoints
- `--max-client-connections` (int) - maximum number of client connections to accept

Example

```
python server.py --port 8080 --host https://0.0.0.0 --timeout
```

This will start a server listening on https://0.0.0.0:8080 and will timeout
after a set number of attempts when no failing to receive a client connection.

---

### Request structure

Requets are sent to the server as a JSON encoded byte string. The JSON object
should look as follows

```
{
TODO
}
```

---

### Model loading

The server does not automatically load a model. It waits for the client to send
a request which contains a model path and name. The server will determine which
version of stable diffusion is in use and which model has been selected
to generate images. It will also determine the best model to load based on
the list of available types in the directory provided.
Modify `settings.py` as you see fit.

---

### Development notes
## Run server and client

See `src/runai/server.py` for an example of how to run the server and `src/runai/client.py` for an example of how to run
the client. Both of these files can be run directly from the command line.

- `StableDiffusionRequestQueueWorker.callback` Handle routes and dispatch to functions
- `socket_server.message_client` Send a message to the client
The socket client will continuously attempt to connect to the server until it is successful. The server will accept
connections from any client on the given port.
2 changes: 1 addition & 1 deletion install.sh
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
#!/usr/bin/bash

pip install -r requirements.txt
pip install -r requirements.txt
5 changes: 0 additions & 5 deletions run.sh
View file

This file was deleted.

3 changes: 2 additions & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
"numpy==1.26.4",
# Core application dependencies
"accelerate==0.29.2",
"huggingface-hub==0.22.2",
"huggingface-hub==0.23.0",
"torch==2.2.2",
"optimum==1.19.1",

Expand All @@ -42,6 +42,7 @@
"pycountry==23.12.11",
"sounddevice==0.4.6", # Required for tts and stt
"pyttsx3==2.90", # Required for tts
"peft==0.12.0",

# Pyinstaller Dependencies
"ninja==1.11.1.1",
Expand Down
21 changes: 21 additions & 0 deletions src/runai/agent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,23 @@ class Agent:
def __init__(self, *args, **kwargs):
self.conversation = kwargs.get("conversation", [])
self.name = kwargs.get("name", "")
self.mood_stats = {
"happy": 0,
"sad": 0,
"neutral": 0,
"angry": 0,
"paranoid": 0,
"anxious": 0,
"excited": 0,
"bored": 0,
"confused": 0,
"relaxed": 0,
"curious": 0,
"frustrated": 0,
"hopeful": 0,
"disappointed": 0,
}
self.user = None

def conversation_so_far(self, use_name=False):
if use_name:
Expand All @@ -15,3 +32,7 @@ def to_dict(self):
"conversation": self.conversation,
"name": self.name
}

def update_mood_stat(self, stat: str, amount: float):
if stat in self.mood_stats:
self.mood_stats[stat] += amount
Loading

0 comments on commit 8f2b1da

Please sign in to comment.