diff --git a/.github/workflows/plugins-docker.yaml b/.github/workflows/plugins-docker.yaml new file mode 100644 index 00000000000..2af37c713b7 --- /dev/null +++ b/.github/workflows/plugins-docker.yaml @@ -0,0 +1,28 @@ +on: + push: + paths: + - "ext/**" + pull_request: + paths: + - "ext/**" + +jobs: + publish-hello-docker-image: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v1 + if: ${{ github.ref_type == 'tag' }} + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: Build and push Docker image + uses: docker/build-push-action@v2 + with: + context: ext/realesrgan + tags: kvalev/realesrgan:0.3.0 + push: true diff --git a/ext/realesrgan/.gitignore b/ext/realesrgan/.gitignore new file mode 100644 index 00000000000..05424f2a4c8 --- /dev/null +++ b/ext/realesrgan/.gitignore @@ -0,0 +1 @@ +weights diff --git a/ext/realesrgan/Dockerfile b/ext/realesrgan/Dockerfile new file mode 100644 index 00000000000..f832c3689d8 --- /dev/null +++ b/ext/realesrgan/Dockerfile @@ -0,0 +1,11 @@ +FROM python:3.10.11-slim-buster + +COPY requirements.txt /app/ +RUN pip install -r /app/requirements.txt + +COPY . /app/ +WORKDIR /app + +EXPOSE 5000 + +CMD ["python3", "main.py"] diff --git a/ext/realesrgan/main.py b/ext/realesrgan/main.py new file mode 100644 index 00000000000..54e2aee4a2e --- /dev/null +++ b/ext/realesrgan/main.py @@ -0,0 +1,85 @@ +import base64 +from io import BytesIO +from basicsr.archs.rrdbnet_arch import RRDBNet +from basicsr.utils.download_util import load_file_from_url +from flask import Flask, request +from realesrgan import RealESRGANer +from realesrgan.archs.srvgg_arch import SRVGGNetCompact +from PIL import Image +import torch + +import logging +import os + +app = Flask(__name__) + +REAL_ESRGAN_MODELS = { + "realesr-general-x4v3": { + "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", + "model_md5": "91a7644643c884ee00737db24e478156", + "scale": 4, + "model": lambda: SRVGGNetCompact( + num_in_ch=3, + num_out_ch=3, + num_feat=64, + num_conv=32, + upscale=4, + act_type="prelu", + ), + }, + "RealESRGAN_x4plus": { + "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", + "model_md5": "99ec365d4afad750833258a1a24f44ca", + "scale": 4, + "model": lambda: RRDBNet( + num_in_ch=3, + num_out_ch=3, + num_feat=64, + num_block=23, + num_grow_ch=32, + scale=4, + ), + }, +} + +name = "realesr-general-x4v3" + +if name not in REAL_ESRGAN_MODELS: + raise ValueError(f"Unknown RealESRGAN model name: {name}") + +model_info = REAL_ESRGAN_MODELS[name] + +model_path = os.path.join('weights', name + '.pth') +if not os.path.isfile(model_path): + ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) + model_path = load_file_from_url(url=model_info["url"], model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None) +logging.info(f"RealESRGAN model path: {model_path}") + +device = "cuda" if torch.cuda.is_available() else "cpu" + +model = RealESRGANer( + scale=model_info["scale"], + model_path=model_path, + model=model_info["model"](), + half=True if "cuda" in device else False, + tile=512, + tile_pad=10, + pre_pad=10, + device=device, +) + +@app.route('/superscale', methods=['POST']) +def superscale(): + scale = request.json.get("scale", model_info["scale"]) + + image_data = base64.b64decode(request.json['image']) + image = Image.open(BytesIO(image_data)).convert('BGR') + + logging.info(f"RealESRGAN input shape: {image.shape}, scale: {scale}") + upsampled = model.enhance(image, outscale=scale)[0] + logging.info(f"RealESRGAN output shape: {upsampled.shape}") + return upsampled + +if __name__ == '__main__': + print("running") + app.run(host='0.0.0.0', port=5000) diff --git a/ext/realesrgan/requirements.txt b/ext/realesrgan/requirements.txt new file mode 100644 index 00000000000..b1a9b1dda23 --- /dev/null +++ b/ext/realesrgan/requirements.txt @@ -0,0 +1,2 @@ +realesrgan==0.3.0 +flask==2.3.2