Duplicate from sergeipetrov/swin2SR-classical-sr-x2-64-IE

handler.py

@@ -0,0 +1,40 @@
+from typing import Dict, List, Any
+from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
+import torch
+import base64
+import logging
+import numpy as np
+from PIL import Image
+from io import BytesIO
+
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+
+# check for GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # load the model
+        self.processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
+        self.model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
+        # move model to device
+        self.model.to(device)
+
+    def __call__(self, data: Any):
+        image = data["inputs"]
+        inputs = self.processor(image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        
+        output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+        output = np.moveaxis(output, source=0, destination=-1)
+        output = (output * 255.0).round().astype(np.uint8)
+
+        img = Image.fromarray(output)
+        buffered = BytesIO()
+        img.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue())
+
+        return img_str.decode()

requirements.txt

@@ -0,0 +1 @@
+pillow