huggingface / optimum-neuron

Easy, fast and very cheap training and inference on AWS Trainium and Inferentia chips.
Apache License 2.0
176 stars 51 forks source link

Make stable diffusion pipelines compatible with compel #581

Closed JingyaHuang closed 2 months ago

JingyaHuang commented 2 months ago

What does this PR do?

Fixes #357

@neo @Suprhimp

Before submitting

HuggingFaceDocBuilderDev commented 2 months ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Suprhimp commented 2 months ago

I tested with my own model even with lora, and it works pretty well.

I compiled with output_hidden_states params true

And I tested with this example code 😃


from optimum.neuron import NeuronStableDiffusionXLPipeline
from compel import Compel, ReturnedEmbeddingsType

pipe = NeuronStableDiffusionXLPipeline.from_pretrained("sdxl_turbo_neuron/", data_parallel_mode="all")
prompt = "Self-portrait oil painting++++++++++, a beautiful cyborg with golden hair, 8k"
negative_prompt = "worst quality"

compel = Compel(
        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
        requires_pooled=[False, True],
    )
prompt_embeds, pooled = compel(prompt)
neg_prompt_embeds, neg_pooled = compel(negative_prompt)

images = pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=neg_prompt_embeds,
        pooled_prompt_embeds=pooled,
        negative_pooled_prompt_embeds=neg_pooled, 
        guidance_scale=0.0, num_inference_steps=1).images

Thanks for the update!

JingyaHuang commented 2 months ago

Thanks @neo and @Suprhimp for the help, the PR is merged. I will include it in our next release!