Merge pull request #57 from aws/add-vision-modality

philschmid · web-flow · commit 1e5771514aff · 2022-02-17T14:17:39.000+01:00
added x-image type
diff --git a/src/sagemaker_huggingface_inference_toolkit/content_types.py b/src/sagemaker_huggingface_inference_toolkit/content_types.py
@@ -26,4 +26,5 @@
 BMP = "image/bmp"
 GIF = "image/gif"
 WEBP = "image/webp"
-VISION_TYPES = [JPEG, PNG, TIFF, BMP, GIF, WEBP]
+X_IMAGE = "image/x-image"
+VISION_TYPES = [JPEG, PNG, TIFF, BMP, GIF, WEBP,X_IMAGE]
diff --git a/src/sagemaker_huggingface_inference_toolkit/decoder_encoder.py b/src/sagemaker_huggingface_inference_toolkit/decoder_encoder.py
@@ -132,6 +132,7 @@ def encode_csv(content):  # type: (str) -> np.array
     content_types.BMP: decode_image,
     content_types.GIF: decode_image,
     content_types.WEBP: decode_image,
+    content_types.X_IMAGE: decode_image,
 }
 
 
diff --git a/tests/integ/test_models_from_hub.py b/tests/integ/test_models_from_hub.py
@@ -65,7 +65,7 @@ def get_all_ecr_images(registry_id, repository_name, result_key):
 @pytest.mark.parametrize(
     "device",
     [
-        "gpu",
+        # "gpu",
         "cpu",
     ],
 )

Original file line number	Diff line number	Diff line change
`@@ -132,6 +132,7 @@ def encode_csv(content): # type: (str) -> np.array`
`132`	`132`	`content_types.BMP: decode_image,`
`133`	`133`	`content_types.GIF: decode_image,`
`134`	`134`	`content_types.WEBP: decode_image,`
	`135`	`+ content_types.X_IMAGE: decode_image,`
`135`	`136`	`}`
`136`	`137`
`137`	`138`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ def get_all_ecr_images(registry_id, repository_name, result_key):`
`65`	`65`	`@pytest.mark.parametrize(`
`66`	`66`	`"device",`
`67`	`67`	`[`
`68`		`- "gpu",`
	`68`	`+ # "gpu",`
`69`	`69`	`"cpu",`
`70`	`70`	`],`
`71`	`71`	`)`