up

Author: yiyixuxu

src/diffusers/models/transformers/transformer_lumina2.py

@@ -241,20 +241,22 @@ def __init__(self, theta: int, axes_dim: List[int], axes_lens: List[int] = (300,
 
     def _precompute_freqs_cis(self, axes_dim: List[int], axes_lens: List[int], theta: int) -> List[torch.Tensor]:
         freqs_cis = []
-        # Use float32 for MPS compatibility
-        dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
         for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
-            emb = get_1d_rotary_pos_embed(d, e, theta=self.theta, freqs_dtype=dtype)
+            emb = get_1d_rotary_pos_embed(d, e, theta=self.theta, freqs_dtype=torch.float64)
             freqs_cis.append(emb)
         return freqs_cis
 
     def _get_freqs_cis(self, ids: torch.Tensor) -> torch.Tensor:
+        device = ids.device
+        if ids.device.type == "mps":
+            ids = ids.to("cpu")
+
         result = []
         for i in range(len(self.axes_dim)):
             freqs = self.freqs_cis[i].to(ids.device)
             index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64)
             result.append(torch.gather(freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
-        return torch.cat(result, dim=-1)
+        return torch.cat(result, dim=-1).to(device)
 
     def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor):
         batch_size = len(hidden_states)

src/diffusers/pipelines/lumina2/pipeline_lumina2.py

@@ -24,8 +24,6 @@
 from ...models.transformers.transformer_lumina2 import Lumina2Transformer2DModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import (
-    is_bs4_available,
-    is_ftfy_available,
     is_torch_xla_available,
     logging,
     replace_example_docstring,
@@ -44,12 +42,6 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-if is_bs4_available():
-    pass
-
-if is_ftfy_available():
-    pass
-
 EXAMPLE_DOC_STRING = """
     Examples:
         ```py