Update handler.py

handler.py

@@ -1,5 +1,5 @@
 from typing import Dict, List, Any
-from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
+from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution, Swin2SRModel
 import torch
 import base64
 import logging
@@ -10,21 +10,20 @@ 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)
+        Swin2SRModel._no_split_modules = ["Swin2SREmbeddings", "Swin2SRStage"]
+        Swin2SRForImageSuperResolution._no_split_modules = ["Swin2SREmbeddings", "Swin2SRStage"]
+        model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64", device_map="auto")
+        model.hf_device_map["upsample"] = model.hf_device_map["swin2sr.embeddings"]
+        self.model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64", device_map=model.hf_device_map)
 
     def __call__(self, data: Any):
         image = data["inputs"]
-        inputs = self.processor(image, return_tensors="pt").to(device)
+        inputs = self.processor(image, return_tensors="pt")
         with torch.no_grad():
             outputs = self.model(**inputs)