Device Placement in diffusers.models.embeddings Rotary Functions

[vroger11]

In diffusers.models.embeddings, functions like get_1d_rotary_pos_embed do not allow users to specify the device of the returned tensors.

Impact

1. Device Mismatch in Stable Audio Open Pipeline
   • In the diffusers.pipelines.stable_audio.pipeline_stable_audio pipeline, this results in some inputs not being on the same device, even after calling pipe.to("cuda").
   • While this does not prevent the pipeline from running, it introduces unnecessary overhead and slightly slows down inference when calling the __call__ method.
2. TensorRT Incompatibility
   • The device inconsistency prevents the transformer model of the stable audio pipeline from being compiled with TensorRT, limiting optimization opportunities.

Suggested Improvement

Allow users to specify the device for tensors returned by rotary embedding functions. Then, modify the call method of the Stable Audio pipeline to use this functionality, ensuring all inputs are consistently placed on the target device and enabling full compatibility with .to("cuda") and TensorRT compilation.

Would love to hear thoughts on this!