Plugin: Upscale plugin based on Real-ESRGAN

.github/workflows/plugins-docker.yaml

@@ -0,0 +1,27 @@
+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
+        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

ext/realesrgan/.gitignore

@@ -0,0 +1 @@
+weights

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"]

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)

ext/realesrgan/requirements.txt

@@ -0,0 +1,2 @@
+realesrgan==0.3.0
+flask==2.3.2