make it compatible to 1 and multiple GPUs

handler.py

@@ -15,22 +15,30 @@ logger.setLevel(logging.DEBUG)
 
 # check for GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
+gpu_count = torch.cuda.device_count()
 
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model        
         self.processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
-        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")
-        logger.info(model.hf_device_map)
-        model.hf_device_map["swin2sr.conv_after_body"] = model.hf_device_map["swin2sr.embeddings"]
-        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)
-
-        print(subprocess.run(["nvidia-smi"]))
-    
+
+        if(gpu_count > 1):
+            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")
+            logger.info(model.hf_device_map)
+            model.hf_device_map["swin2sr.conv_after_body"] = model.hf_device_map["swin2sr.embeddings"]
+            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)
+
+            print(subprocess.run(["nvidia-smi"]))
+
+        else:
+            self.model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
+            # move model to device
+            self.model.to(device)
+
+
     def __call__(self, data: Any):
         """
         Args:
@@ -41,7 +49,11 @@ class EndpointHandler:
         """
 
         image = data["inputs"]
-        inputs = self.processor(image, return_tensors="pt")
+
+        if(gpu_count > 1):
+            inputs = self.processor(image, return_tensors="pt")
+        else:
+            inputs = self.processor(image, return_tensors="pt").to(device)
 
         try:
             with torch.no_grad():