Add support to run as Docker Container

.dockerignore

@@ -0,0 +1,28 @@
+__pycache__
+.ruff_cache
+/cache.json
+/*.json
+*.log
+*.log.*
+*.bak
+*.ckpt
+*.safetensors
+*.pth
+*.pt
+*.bin
+*.optim
+*.lock
+*.zip
+*.rar
+*.7z
+*.pyc
+/*.bat
+/*.sh
+/*.mp3
+/*.lnk
+!webui.bat
+!webui.sh
+!package.json
+!config.json
+!ui-config.json
+!params.txt

README.md

@@ -298,7 +298,7 @@ check [ChangeLog](CHANGELOG.md) for when feature was first introduced as it will
 ### Sponsors
 
 <div align="center">
-<!-- sponsors --><a href="https://github.com/allangrant"><img src="https://github.com/allangrant.png" width="60px" alt="Allan Grant" /></a><a href="https://github.com/BrentOzar"><img src="https://github.com/BrentOzar.png" width="60px" alt="Brent Ozar" /></a><a href="https://github.com/inktomi"><img src="https://github.com/inktomi.png" width="60px" alt="Matthew Runo" /></a><a href="https://github.com/4joeknight4"><img src="https://github.com/4joeknight4.png" width="60px" alt="" /></a><a href="https://github.com/mantzaris"><img src="https://github.com/mantzaris.png" width="60px" alt="a.v.mantzaris" /></a><a href="https://github.com/CurseWave"><img src="https://github.com/CurseWave.png" width="60px" alt="" /></a><a href="https://github.com/freeload101"><img src="https://github.com/freeload101.png" width="60px" alt="Robert McCurdy" /></a><!-- sponsors -->
+<!-- sponsors --><a href="https://github.com/allangrant"><img src="https://github.com/allangrant.png" width="60px" alt="Allan Grant" /></a><a href="https://github.com/BrentOzar"><img src="https://github.com/BrentOzar.png" width="60px" alt="Brent Ozar" /></a><a href="https://github.com/inktomi"><img src="https://github.com/inktomi.png" width="60px" alt="Matthew Runo" /></a><a href="https://github.com/4joeknight4"><img src="https://github.com/4joeknight4.png" width="60px" alt="" /></a><a href="https://github.com/mantzaris"><img src="https://github.com/mantzaris.png" width="60px" alt="a.v.mantzaris" /></a><a href="https://github.com/CurseWave"><img src="https://github.com/CurseWave.png" width="60px" alt="" /></a><!-- sponsors -->
 </div>
 
 <br>

docker/Dockerfile

@@ -0,0 +1,13 @@
+ARG BASE_IMG
+FROM ${BASE_IMG}
+
+USER root
+WORKDIR /workspace
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+COPY ./ ./
+
+RUN ./docker/setup.sh
+
+ENTRYPOINT ["./docker/entrypoint.sh"]

docker/container.py

@@ -0,0 +1,43 @@
+from helper import cmd, IMG, args
+
+name, port, compute, data_dir, Vol, noVol = args
+
+# Setup base args
+CMD = ["./webui.sh", "--uv", "--listen"]
+dockerArgs = ["-t", f"-p {port}:7860", f"--name {name}", '-e "venv_dir=/python/venv"']
+
+# Setup data dir
+if data_dir:
+ dockerArgs.append(f'-e "SD_DATADIR={data_dir}"')
+ Vol["Data"] = data_dir
+Vol = Vol.items()
+
+# Setup compute platform args
+if compute == "cuda":
+ dockerArgs.append("--gpus all")
+ CMD.append("--use-cuda")
+elif compute == "rocm":
+ dockerArgs.extend(["--device /dev/dri", "--device /dev/kfd"])
+ CMD.append("--use-rocm")
+else: CMD.append("--use-cpu=all")
+
+# Setup volume args
+if not noVol:
+ for vName, vPath in Vol:
+ dockerArgs.append(f'-v "SD-Next_{vName}:{vPath}"')
+
+dockerArgs= " ".join(dockerArgs)
+CMD= " ".join(CMD)
+print(f'''
+Container Settings:
+ Name: {name}
+ Port: {port}
+ Compute Platform: {compute}
+ Volumes: {"Disabled" if noVol else ", ".join([f'{key} -> {value}' for key, value in Vol])}
+ Docker Args: {dockerArgs}
+ Container CMD: {CMD}
+''')
+cmd(f"docker container rm {name} -f", msg=f"Removing container named {name}...")
+cmd(f"docker image rm sd-next -f", msg="Removing image named sd-next...")
+cmd(f'docker build -t sd-next -f ./Dockerfile --build-arg="BASE_IMG={IMG[compute]}" ../', msg="Building Docker Image (might takes few minutes)...")
+cmd(f'docker run {dockerArgs} sd-next {CMD}', msg="Running SD Next Container...", stream=True)

docker/entrypoint.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+
+cd "$(dirname "$(dirname "$0")")"
+
+echo
+echo "======================="
+echo "== SD Next Container =="
+echo "======================="
+echo
+
+if [ $# -eq 0 ]; then
+ exec /bin/bash
+else
+ exec "$@"
+fi

docker/helper.py

@@ -0,0 +1,96 @@
+import argparse
+import subprocess
+import os
+import re
+
+IMG = {
+ "cuda": "nvidia/cuda:12.1.1-runtime-ubuntu22.04",
+ "rocm": "rocm/dev-ubuntu-22.04:6.0.2",
+ "cpu": "ubuntu:22.04",
+}
+
+# Arg parser utils
+class MultilineHelpFormatter(argparse.HelpFormatter):
+ def __init__(self, prog):
+ super().__init__("", 2, 55, 200)
+ def _split_lines(self, text, _):
+ lines = text.splitlines()
+ lines.append("")
+ return lines
+class VolumeAction(argparse.Action):
+ def __call__(self, parser, namespace, values, option_string=None):
+ dictionary = getattr(namespace, self.dest)
+ data_dir = getattr(namespace, "data_dir")
+
+ if ":" in values:
+ key, value = values.split(':', 1)
+ if data_dir and key == "Data": raise argparse.ArgumentTypeError('Volume "Data" is reserved, as --data-dir presents')
+ if value:
+ dictionary[key] = value
+ elif key in dictionary:
+ del dictionary[key]
+ else:
+ print(f"Invalid item '{values}'. Items must be in the format 'VOL_NAME:VOL_PATH'")
+def nonEmptyString(value):
+ if value == "":
+ raise argparse.ArgumentTypeError("Cannot be empty")
+ return value
+
+# Initialize and expose args
+argParser = argparse.ArgumentParser(conflict_handler='resolve', add_help=True, formatter_class=MultilineHelpFormatter)
+argParser.add_argument('-n', '--name', type=str, default = os.environ.get("SD_CONTAINER_NAME","SD-Next"), help = '''\
+Specify the name for the container
+Default: SD-Next
+''')
+argParser.add_argument('-p', '--port', type=int, default = os.environ.get("SD_PORT",7860), help = '''\
+Specify the port exposed by the container
+Default: 7860
+''')
+argParser.add_argument('--compute', type=str, choices=['cuda', 'rocm', 'cpu'], default = os.environ.get("SD_CONTAINER_COMPUTE","cuda"), help = '''\
+Specify the compute platform use by the container
+Default: cuda
+''')
+argParser.add_argument('--data-dir', type=nonEmptyString, default = os.environ.get("SD_DATADIR",None), help = '''\
+Specify the directory for SD Next data
+Default: None
+''')
+argParser.add_argument('-v', '--volumes', action=VolumeAction, default={
+ "Cache": "/root/.cache",
+ "Python": "/python"
+}, metavar='VOL_NAME:VOL_PATH', help='''\
+Mount a volume to the container
+This flag can be used multiple times for multiple volumes mount
+If you want to remove default volume, type "-v DEFAULT_VOL_NAME:" (leave VOL_PATH as empty)
+Default:
+ Data:[Path Specified With --data-dir] (save SD Next data to volume)
+ Reserved for --data-dir. Unless --data-dir is not provided, it will always set to --data-dir
+ Python:/python (save the venv to volume - venv will be created at /python/venv)
+ Cache:/root/.cache (save the cache generated by pip, uv, huggingface)
+'''
+)
+argParser.add_argument('--no-volume', default = os.environ.get("SD_CONTAINER_NO_VOL",False), action='store_true', help = '''\
+Disable volume mounting (including default volume)
+Default: False
+''')
+args, _ = argParser.parse_known_args()
+args = vars(args).values()
+
+# Command runner
+wd = os.path.dirname(os.path.abspath(__file__))
+log = open(os.path.join(wd, "../docker.log"), "a+")
+log.truncate(0)
+def cmd(cmd, msg=None, stream=False):
+ print(msg or f"Running - {cmd}\n")
+ process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, bufsize=1, cwd=wd)
+ for line in iter(process.stdout.readline, ''):
+ if stream: print(line, end='')
+ ansi_escape = re.compile(r'\x1b[^m]*m')
+ line = ansi_escape.sub('', line)
+ log.write(line)
+ log.flush()
+ log.write("")
+ process.wait()
+ if process.returncode != 0:
+ print("An error has occurred, check docker.log for the details")
+ exit(1)
+ print("")

docker/setup.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+echo
+echo ======================
+echo == Updating apt-get ==
+echo ======================
+echo
+apt-get update -y
+
+echo
+echo =========================
+echo == Setting up apt-fast ==
+echo =========================
+echo
+apt-get install curl wget aria2 -y
+/bin/bash -c "$(curl -sL https://git.io/vokNn)"
+echo -e "_MAXNUM=16\n_MAXCONPERSRV=16\n_SPLITCON=64" > /etc/apt-fast.conf
+
+echo
+echo =========================
+echo == Installing packages ==
+echo =========================
+echo
+apt-fast install -y git \
+ python3.10 pythonpy python3.10-venv python3-pip \
+ curl wget aria2