Add diffusers utils

[philschmid]

What does this PR do?

This PR adds support for diffusers and text-to-image task for zero-code deployments. This will allow customers to deploy any diffusers model support by AutoPipeline including:

• Stable Diffusion
• ControlNet
• Stable Diffusion XL (SDXL)
• DeepFloyd IF
• Kandinsky
• Kandinsky 2.2

This PR adds a new:

• diffusers_utils.py which handles the heavy lifting for the zero-code deployment with new unit and integration tests
• additional decoder for serializing the generated image into a binary for response
• support for loading safetensors and sharded model from the hugging faced hub.

Test locally

1. manually change MMS_CONFIG_FILE

wget -O sagemaker-mms.properties https://github.com/raw/aws/deep-learning-containers/master/huggingface/build_artifacts/inference/config.properties

2. Run Container, e.g. text-to-image

HF_MODEL_ID="stabilityai/stable-diffusion-xl-base-1.0" HF_TASK="text-to-image" python src/sagemaker_huggingface_inference_toolkit/serving.py

3. Adjust handler_service.py and comment out if content_type in content_types.UTF8_TYPES: thats needed for SageMaker but cannot be used locally

4. Send request

curl --request POST \
  --url http://localhost:8080/invocations \
  --header 'Accept: image/png' \
  --header 'Content-Type: application/json' \
  --data '"{\"inputs\": \"Camera\"}" \
  --output image.png

[davidthomas426]

From your description:
"3. Adjust handler_service.py and comment out if content_type in content_types.UTF8_TYPES: thats needed for SageMaker but cannot be used locally"

Why is that needed for SageMaker but cannot be used locally? That's a strange requirement for testing locally and does not give me confidence that everything here works correctly.

[philschmid]

From your description: "3. Adjust handler_service.py and comment out if content_type in content_types.UTF8_TYPES: thats needed for SageMaker but cannot be used locally"

Why is that needed for SageMaker but cannot be used locally? That's a strange requirement for testing locally and does not give me confidence that everything here works correctly.

It seems that when the container is deployed on SageMaker, the "body" is stringified and encoded. When starting the MMS server locally (no container), the body is just what you pass. Since we are not including the dockerfile here, we need to use MMS locally.

[davidthomas426]

From your description: "3. Adjust handler_service.py and comment out if content_type in content_types.UTF8_TYPES: thats needed for SageMaker but cannot be used locally"
Why is that needed for SageMaker but cannot be used locally? That's a strange requirement for testing locally and does not give me confidence that everything here works correctly.

It seems that when the container is deployed on SageMaker, the "body" is stringified and encoded. When starting the MMS server locally (no container), the body is just what you pass. Since we are not including the dockerfile here, we need to use MMS locally.

Ok, but for whatever reason the "body" is a bytes or a str when called, your code could just handle that properly. So, if it's already a str, don't call decode. This is just python3 unicode stuff.

json.loads will call decode("utf-8") for you if handed a bytes. See the below:

>>> json_string = '{"🙃": 17}'
>>> json_bytes = json_string.encode("utf-8")

>>> dump = lambda v: print(f"{type(v)}  =>  {repr(v)}")
>>> dump(json_string)
<class 'str'>  =>  '{"🙃": 17}'
>>> dump(json_bytes)
<class 'bytes'>  =>  b'{"\xf0\x9f\x99\x83": 17}'

>>> json.loads(json_string)
{'🙃': 17}
>>> json.loads(json_bytes)
{'🙃': 17}

[davidthomas426]

Would it makes sense to go to else case if task == "text-to-image" but is_diffusers_availablereturnsFalse`?

[philschmid]

no. There you need both.

[davidthomas426]

But that's what this if-elif-else chain will do. If is_diffusers_available() is False and task is "text-to-image", then this branch will fail and it will go to the else clause.

So you're saying that doesn't make sense?

[philschmid]

You mean we should rather error? That when task-to-image == True but diffusers not available?

[davidthomas426]

Yeah, probably. Right now, the code does something that doesn't make sense in that case, because it falls through to the else clause.

I said that from the beginning. But yeah, at this point this back and forth is a bit ridiculous, so whatever. It's probably fine in practice. But still, this if-elif-else logic seems kind of wrong. I don't know how to be more clear.

[davidthomas426]

This if-elif-else seems odd to me. Could the if ever match when we want to use diffusers?

Also, why is this extra logic in the if case regarding bfloat16 and whatever else only used when TRUST_REMOTE_CODE is True, but then TRUST_REMOTE_CODE is also passed in other cases?

[philschmid]

What is odd for you?

1. branch is when you have TRUST_REMOTE_CODE which is for custom modelling available through the hub and it only work if there is a HF_MODEL_ID and you have GPUs. -> was needed for mpt.
2. checks whether we have diffusers installed and if we have a diffusers task and to load image generation model.
3. is default where we try to load the custom pipeline

[davidthomas426]

It's strange because there's special bfloat16 logic but only when TRUST_REMOTE_CODE is True. I mentioned that in my comment. And it also seemed strange to me because of what I said in my other comment below.

[philschmid]

Ok, but for whatever reason the "body" is a bytes or a str when called, your code could just handle that properly. So, if it's already a str, don't call decode. This is just python3 unicode stuff.

json.loads will call decode("utf-8") for you if handed a bytes. See the below:

Thats what was implemented and what SageMaker is doing. I don't want to change something which is not wrong. It is just for testing. Or can you guarantee its not needed?

[davidthomas426]

I'm approving because I don't want to hold things up. But the lack of bfloat16 logic in any branch but the first TRUST_REMOTE_CODE one, and the other issue I pointed out with the if-elif-else clause, seem kind of wrong to me.

In practice, it may be fine. I'll leave that up to you.

[davidthomas426]

Yeah, probably. Right now, the code does something that doesn't make sense in that case, because it falls through to the else clause.

I said that from the beginning. But yeah, at this point this back and forth is a bit ridiculous, so whatever. It's probably fine in practice. But still, this if-elif-else logic seems kind of wrong. I don't know how to be more clear.

[philschmid]

I'm approving because I don't want to hold things up. But the lack of bfloat16 logic in any branch but the first TRUST_REMOTE_CODE one, and the other issue I pointed out with the if-elif-else clause, seem kind of wrong to me.

Let me double check and make sure the tests are green.