From 8f1dc00878bae55f498fd4f5b7596c233b62839e Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Apr 2025 14:51:04 +0000 Subject: [PATCH 1/6] Created more generic pipeline for text-to-image task --- .../diffusers_utils.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py index 068a41a..7d4fcfb 100644 --- a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py +++ b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py @@ -29,24 +29,22 @@ def is_diffusers_available(): if is_diffusers_available(): import torch - from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler, StableDiffusionPipeline + from diffusers import DiffusionPipeline -class SMAutoPipelineForText2Image: +class DiffusionPipelineForText2Image: + def __init__(self, model_dir: str, device: str = None): # needs "cuda" for GPU + self.pipeline = None dtype = torch.float32 if device == "cuda": dtype = torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float16 - device_map = "auto" if device == "cuda" else None + if torch.cuda.device_count() > 1: + device_map = "balanced" + self.pipeline = DiffusionPipeline.from_pretrained(model_dir, torch_dtype=dtype, device_map=device_map) - self.pipeline = AutoPipelineForText2Image.from_pretrained(model_dir, torch_dtype=dtype, device_map=device_map) - # try to use DPMSolverMultistepScheduler - if isinstance(self.pipeline, StableDiffusionPipeline): - try: - self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config) - except Exception: - pass - self.pipeline.to(device) + if not self.pipeline: + self.pipeline = DiffusionPipeline.from_pretrained(model_dir, torch_dtype=dtype).to(device) def __call__( self, @@ -64,7 +62,7 @@ def __call__( DIFFUSERS_TASKS = { - "text-to-image": SMAutoPipelineForText2Image, + "text-to-image": DiffusionPipelineForText2Image, } From dd82b41b847177035b26f48a945a2a9ca224672a Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Apr 2025 15:09:51 +0000 Subject: [PATCH 2/6] Updated test --- tests/unit/test_diffusers_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_diffusers_utils.py b/tests/unit/test_diffusers_utils.py index c00c139..83c6d2d 100644 --- a/tests/unit/test_diffusers_utils.py +++ b/tests/unit/test_diffusers_utils.py @@ -16,7 +16,7 @@ from transformers.testing_utils import require_torch, slow from PIL import Image -from sagemaker_huggingface_inference_toolkit.diffusers_utils import SMAutoPipelineForText2Image +from sagemaker_huggingface_inference_toolkit.diffusers_utils import DiffusionPipelineForText2Image from sagemaker_huggingface_inference_toolkit.transformers_utils import _load_model_from_hub, get_pipeline From d1de78ec8c2ae216d38d789fc7e27b17d7fae327 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Apr 2025 15:11:18 +0000 Subject: [PATCH 3/6] black --- src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py index 7d4fcfb..75b4357 100644 --- a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py +++ b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py @@ -33,7 +33,7 @@ def is_diffusers_available(): class DiffusionPipelineForText2Image: - + def __init__(self, model_dir: str, device: str = None): # needs "cuda" for GPU self.pipeline = None dtype = torch.float32 From de77f7eec90dad89177366e4850427284487a85b Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Fri, 11 Apr 2025 15:11:48 +0000 Subject: [PATCH 4/6] Fix test --- tests/unit/test_diffusers_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_diffusers_utils.py b/tests/unit/test_diffusers_utils.py index 83c6d2d..68cb3cb 100644 --- a/tests/unit/test_diffusers_utils.py +++ b/tests/unit/test_diffusers_utils.py @@ -28,7 +28,7 @@ def test_get_diffusers_pipeline(): tmpdirname, ) pipe = get_pipeline("text-to-image", -1, storage_dir) - assert isinstance(pipe, SMAutoPipelineForText2Image) + assert isinstance(pipe, DiffusionPipelineForText2Image) @slow From e362fa0af29bd15745e44ce67a347e352528146e Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Mon, 14 Apr 2025 14:38:48 +0000 Subject: [PATCH 5/6] Renamed pipeline --- .../diffusers_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py index 75b4357..c0de779 100644 --- a/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py +++ b/src/sagemaker_huggingface_inference_toolkit/diffusers_utils.py @@ -32,7 +32,7 @@ def is_diffusers_available(): from diffusers import DiffusionPipeline -class DiffusionPipelineForText2Image: +class SMDiffusionPipelineForText2Image: def __init__(self, model_dir: str, device: str = None): # needs "cuda" for GPU self.pipeline = None @@ -62,7 +62,7 @@ def __call__( DIFFUSERS_TASKS = { - "text-to-image": DiffusionPipelineForText2Image, + "text-to-image": SMDiffusionPipelineForText2Image, } From 9797ee8e5c06905345f9ef5c1867ea9d944e2dd9 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Mon, 14 Apr 2025 14:40:42 +0000 Subject: [PATCH 6/6] Updated test with new pipeline name --- tests/unit/test_diffusers_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_diffusers_utils.py b/tests/unit/test_diffusers_utils.py index 68cb3cb..52a7929 100644 --- a/tests/unit/test_diffusers_utils.py +++ b/tests/unit/test_diffusers_utils.py @@ -16,7 +16,7 @@ from transformers.testing_utils import require_torch, slow from PIL import Image -from sagemaker_huggingface_inference_toolkit.diffusers_utils import DiffusionPipelineForText2Image +from sagemaker_huggingface_inference_toolkit.diffusers_utils import SMDiffusionPipelineForText2Image from sagemaker_huggingface_inference_toolkit.transformers_utils import _load_model_from_hub, get_pipeline @@ -28,7 +28,7 @@ def test_get_diffusers_pipeline(): tmpdirname, ) pipe = get_pipeline("text-to-image", -1, storage_dir) - assert isinstance(pipe, DiffusionPipelineForText2Image) + assert isinstance(pipe, SMDiffusionPipelineForText2Image) @slow