sergeipetrov
/

swin2SR-classical-sr-x2-64-IE-distr

Inference Endpoints

Model card Files Files and versions Community

sergeipetrov commited on Jun 26

Commit

686892f

•

1 Parent(s): 92c4446

Update handler.py

Files changed (1) hide show

handler.py +1 -0

handler.py CHANGED Viewed

@@ -18,6 +18,7 @@ class EndpointHandler:
         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["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)

         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)