allow usage of multiple GPUs

details see https://github.com/pimcore-labs/hugging-face-eap/issues/21#issuecomment-2200537717

handler.py

@@ -16,11 +16,15 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # load the model
+        # 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")
+        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)
 
     def __call__(self, data: Any):
         """
@@ -32,7 +36,7 @@ class EndpointHandler:
         """
 
         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)