AttributeError in `diffusers.models.transformers.transformer_flux` - Missing `rope` Attribute

[rajveer43]

Describe the bug

I encountered an issue while working with the diffusers library where an AttributeError is raised because the transformer_flux module does not have an attribute named rope. This problem arises when trying to modify the rope function to handle CUDA devices.

Reproduction

Below is a code snippet that leads to this error:

import torch
from diffusers import FluxPipeline
import diffusers
from PIL import Image
import matplotlib.pyplot as plt

# Attempt to modify the rope function to handle CUDA device
_flux_rope = diffusers.models.transformers.transformer_flux.rope
def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
    assert dim % 2 == 0, "The dimension must be even."
    if pos.device.type == "cuda":
        # Move tensor to CPU for ROPE computation, then move it back to CUDA
        return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device)
    else:
        # Perform ROPE computation directly if tensor is not on CUDA
        return _flux_rope(pos, dim, theta)

# Replace the original rope function with the new one
diffusers.models.transformers.transformer_flux.rope = new_flux_rope

The rope function should be accessible in the transformer_flux module to allow customization, such as handling CUDA devices or modifying ROPE computations.

Logs

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-17-f5d61243cd5c> in <cell line: 8>()
      6 
      7 # Modify the rope function to handle the CPU device only
----> 8 _flux_rope = diffusers.models.transformers.transformer_flux.rope
      9 def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
     10     assert dim % 2 == 0, "The dimension must be even."

AttributeError: module 'diffusers.models.transformers.transformer_flux' has no attribute 'rope'

System Info

0.31.0.dev0(diffuser version)
3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0](python iversion)

Who can help?

@sayakpaul @DN6

[yiyixuxu]

hi @rajveer43
we did a refactor and replaced the rope method with the get_1d_rotary_pos_embed see this PR #9074

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

Closing due to inactivity.