[Bugfix][Ascend NPU][Boogu] Replace complex RoPE with real-valued operations (#6571)

Signed-off-by: lea_lei <JiaWen_Lei_7@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -23,9 +23,13 @@ NORM_EPS = 1e-5
 
 
 @pytest.fixture(autouse=True)
-def _init_distributed():
+def _init_distributed(request):
     """Initialize the minimal single-rank distributed environment required by
     the vLLM parallel linear layers (tensor-parallel group must exist)."""
+    if request.node.name.startswith("test_real_rotary_"):
+        yield
+        return
+
     from vllm.distributed.parallel_state import (
         cleanup_dist_env_and_memory,
         init_distributed_environment,
@@ -70,11 +74,11 @@ def _randomize_parameters(module: torch.nn.Module) -> None:
             param.uniform_(-0.02, 0.02)
 
 
-def _identity_rotary_emb(batch_size: int, seq_len: int) -> torch.Tensor:
-    """Complex rotary frequencies encoding a zero rotation."""
-    return torch.polar(
-        torch.ones(batch_size, seq_len, HEAD_DIM // 2),
-        torch.zeros(batch_size, seq_len, HEAD_DIM // 2),
+def _identity_rotary_emb(batch_size: int, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Real-valued rotary frequencies encoding a zero rotation."""
+    return (
+        torch.ones(batch_size, seq_len, HEAD_DIM),
+        torch.zeros(batch_size, seq_len, HEAD_DIM),
     )
 
 
@@ -119,6 +123,72 @@ def test_boogu_image_transformer_import():
     assert BooguImageTransformer2DModel is not None
 
 
+def test_real_rotary_emb_matches_complex_reference():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import apply_rotary_emb
+
+    torch.manual_seed(0)
+    x = torch.randn(2, 7, NUM_HEADS, HEAD_DIM)
+    angles = torch.randn(2, 7, HEAD_DIM // 2)
+    rotary_emb = (
+        angles.cos().repeat_interleave(2, dim=-1),
+        angles.sin().repeat_interleave(2, dim=-1),
+    )
+
+    actual = apply_rotary_emb(x, rotary_emb)
+    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], HEAD_DIM // 2, 2))
+    freqs_cis = torch.polar(torch.ones_like(angles), angles).unsqueeze(2)
+    expected = torch.view_as_real(x_complex * freqs_cis).flatten(3).type_as(x)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_real_rotary_frequency_tables_are_real_and_repeated():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamRotaryPosEmbed,
+    )
+
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(AXES_DIM_ROPE, AXES_LENS, theta=10000)
+
+    assert len(freqs_real) == len(AXES_DIM_ROPE)
+    for (freqs_cos, freqs_sin), axis_dim, axis_len in zip(freqs_real, AXES_DIM_ROPE, AXES_LENS):
+        assert freqs_cos.shape == (axis_len, axis_dim)
+        assert freqs_sin.shape == (axis_len, axis_dim)
+        assert not freqs_cos.is_complex()
+        assert not freqs_sin.is_complex()
+        torch.testing.assert_close(freqs_cos[..., ::2], freqs_cos[..., 1::2])
+        torch.testing.assert_close(freqs_sin[..., ::2], freqs_sin[..., 1::2])
+
+
+def test_real_rotary_frequency_tables_use_float32_on_npu(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
+
+    npu_platform = SimpleNamespace(is_npu=lambda: True, supports_float64=lambda: True)
+    monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", npu_platform)
+
+    freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+        AXES_DIM_ROPE, AXES_LENS, theta=10000
+    )
+
+    assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
+
+
+def test_real_rotary_frequency_tables_use_float32_on_mps(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
+
+    platform = SimpleNamespace(
+        is_npu=lambda: False,
+        supports_float64=lambda: True,
+    )
+    monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", platform)
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+
+    freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+        AXES_DIM_ROPE, AXES_LENS, theta=10000
+    )
+
+    assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
+
+
 def test_single_stream_block_shape():
     from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
         BooguImageTransformerBlock,
@@ -398,10 +468,10 @@ def test_transformer_forward_t2i_shape():
     timestep = torch.full((batch_size,), 0.5)
     instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
     instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
-    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000)
 
     with torch.no_grad():
-        out = model(latents, timestep, instruction_hidden_states, freqs_cis, instruction_attention_mask)
+        out = model(latents, timestep, instruction_hidden_states, freqs_real, instruction_attention_mask)
 
     assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
     assert torch.isfinite(out).all()
@@ -431,7 +501,7 @@ def test_transformer_forward_ti2i_shape():
     timestep = torch.full((batch_size,), 0.5)
     instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
     instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
-    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000)
 
     # One sample, one reference image (Boogu editing supports a single ref).
     ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]]
@@ -441,7 +511,7 @@ def test_transformer_forward_ti2i_shape():
             latents,
             timestep,
             instruction_hidden_states,
-            freqs_cis,
+            freqs_real,
             instruction_attention_mask,
             ref_image_hidden_states=ref_image_hidden_states,
         )

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -467,7 +467,8 @@ class _FakeTransformer:
         "reduce_type": "mean",
     }
 
-    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+        assert all(not freqs.is_complex() for pair in freqs_real for freqs in pair)
         return torch.zeros_like(latents)
 
 
@@ -812,7 +813,7 @@ class _RecordingRefTransformer(_FakeTransformer):
     def __init__(self):
         self.calls = []
 
-    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
         self.calls.append(kwargs.get("ref_image_hidden_states"))
         return torch.zeros_like(latents)
 

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -34,20 +34,33 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+RotaryEmbedding = tuple[torch.Tensor, torch.Tensor]
+RotaryFrequencyTables = list[RotaryEmbedding]
 
-def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
-    """Apply rotary embeddings in complex form (upstream `use_real=False` path).
+
+def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tensor:
+    """Apply rotary embeddings with real-valued cosine and sine tensors.
 
     Args:
         x: Query or key tensor of shape [B, S, H, D].
-        freqs_cis: Complex frequency tensor of shape [B, S, D // 2].
+        rotary_emb: Cosine and sine tensors, each shaped [B, S, D].
     """
-    x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2))
-    freqs_cis = freqs_cis.unsqueeze(2)
-    x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+    freqs_cos, freqs_sin = rotary_emb
+    freqs_cos = freqs_cos.unsqueeze(2)
+    freqs_sin = freqs_sin.unsqueeze(2)
+
+    # The frequency values are repeated for each adjacent channel pair. This
+    # is the real-valued equivalent of multiplying view_as_complex(x) by cis.
+    x_float = x.float()
+    x_even = x_float[..., ::2]
+    x_odd = x_float[..., 1::2]
+    cos = freqs_cos[..., ::2]
+    sin = freqs_sin[..., ::2]
+    x_out = torch.stack((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1).flatten(3)
     return x_out.type_as(x)
 
 
@@ -173,7 +186,7 @@ class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
 
 
 class BooguImageDoubleStreamRotaryPosEmbed(nn.Module):
-    """3-axis RoPE producing per-segment frequency tensors.
+    """3-axis RoPE producing real-valued per-segment embeddings.
 
     Returns separate frequency tensors for instruction tokens, reference-image
     tokens, noise-image tokens, the full joint sequence, and the combined
@@ -194,31 +207,54 @@ class BooguImageDoubleStreamRotaryPosEmbed(nn.Module):
         self.patch_size = patch_size
 
     @staticmethod
-    def get_freqs_cis(
+    def get_freqs_real(
         axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
-    ) -> list[torch.Tensor]:
-        freqs_cis = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+    ) -> RotaryFrequencyTables:
+        freqs_real = []
+        # Ascend supports real-valued gather but float32 avoids routing the
+        # entire RoPE path through slower and less broadly supported float64
+        # kernels. MPS also requires float32 because it does not support
+        # float64. Other platforms retain the previous precision where valid.
+        freqs_dtype = (
+            torch.float32
+            if (
+                current_omni_platform.is_npu()
+                or not current_omni_platform.supports_float64()
+                or torch.backends.mps.is_available()
+            )
+            else torch.float64
+        )
         for d, e in zip(axes_dim, axes_lens):
-            emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
-            freqs_cis.append(emb)
-        return freqs_cis
+            freqs_cos, freqs_sin = get_1d_rotary_pos_embed(
+                d,
+                e,
+                theta=theta,
+                use_real=True,
+                repeat_interleave_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            freqs_real.append((freqs_cos, freqs_sin))
+        return freqs_real
 
-    def _get_freqs_cis(self, freqs_cis: list[torch.Tensor], ids: torch.Tensor) -> torch.Tensor:
+    def _get_freqs_real(self, freqs_real: RotaryFrequencyTables, ids: torch.Tensor) -> RotaryEmbedding:
         device = ids.device
         if ids.device.type == "mps":
             ids = ids.to("cpu")
 
-        result = []
+        cos_result = []
+        sin_result = []
         for i in range(len(self.axes_dim)):
-            freqs = 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).to(device)
+            freqs_cos, freqs_sin = freqs_real[i]
+            freqs_cos = freqs_cos.to(ids.device)
+            freqs_sin = freqs_sin.to(ids.device)
+            index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64)
+            cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+            sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+        return torch.cat(cos_result, dim=-1).to(device), torch.cat(sin_result, dim=-1).to(device)
 
     def forward(
         self,
-        freqs_cis: list[torch.Tensor],
+        freqs_real: RotaryFrequencyTables,
         attention_mask: torch.Tensor,
         l_effective_ref_img_len: list[list[int]],
         l_effective_img_len: list[int],
@@ -286,47 +322,52 @@ class BooguImageDoubleStreamRotaryPosEmbed(nn.Module):
             position_ids[i, pe_shift_len:seq_len, 1] = row_ids
             position_ids[i, pe_shift_len:seq_len, 2] = col_ids
 
-        freqs_cis = self._get_freqs_cis(freqs_cis, position_ids)
+        freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids)
 
-        cap_freqs_cis = torch.zeros(
-            batch_size, encoder_seq_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
-        )
-        ref_img_freqs_cis = torch.zeros(
-            batch_size, max_ref_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
-        )
-        img_freqs_cis = torch.zeros(batch_size, max_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype)
+        def empty_rotary_emb(seq_len: int) -> RotaryEmbedding:
+            shape = (batch_size, seq_len, freqs_cos.shape[-1])
+            return (
+                torch.zeros(shape, device=device, dtype=freqs_cos.dtype),
+                torch.zeros(shape, device=device, dtype=freqs_sin.dtype),
+            )
+
+        cap_rotary_emb = empty_rotary_emb(encoder_seq_len)
+        ref_img_rotary_emb = empty_rotary_emb(max_ref_img_len)
+        img_rotary_emb = empty_rotary_emb(max_img_len)
 
         combined_img_seq_lengths = [
             sum(ref_img_len) + img_len for ref_img_len, img_len in zip(l_effective_ref_img_len, l_effective_img_len)
         ]
         max_combined_img_len = max(combined_img_seq_lengths)
 
-        combined_img_freqs_cis = torch.zeros(
-            batch_size, max_combined_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
-        )
+        combined_img_rotary_emb = empty_rotary_emb(max_combined_img_len)
 
         for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
             zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)
         ):
-            cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len]
-            ref_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
-            img_freqs_cis[i, :img_len] = freqs_cis[
-                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
-            ]
-
-            combined_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
-            combined_img_freqs_cis[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = freqs_cis[
-                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
-            ]
+            ref_len = sum(ref_img_len)
+            img_start = cap_seq_len + ref_len
+            for source, cap_target, ref_target, img_target, combined_target in zip(
+                (freqs_cos, freqs_sin),
+                cap_rotary_emb,
+                ref_img_rotary_emb,
+                img_rotary_emb,
+                combined_img_rotary_emb,
+            ):
+                cap_target[i, :cap_seq_len] = source[i, :cap_seq_len]
+                ref_target[i, :ref_len] = source[i, cap_seq_len:img_start]
+                img_target[i, :img_len] = source[i, img_start : img_start + img_len]
+                combined_target[i, :ref_len] = source[i, cap_seq_len:img_start]
+                combined_target[i, ref_len : ref_len + img_len] = source[i, img_start : img_start + img_len]
 
         return (
-            cap_freqs_cis,
-            ref_img_freqs_cis,
-            img_freqs_cis,
-            freqs_cis,
+            cap_rotary_emb,
+            ref_img_rotary_emb,
+            img_rotary_emb,
+            (freqs_cos, freqs_sin),
             l_effective_cap_len,
             seq_lengths,
-            combined_img_freqs_cis,
+            combined_img_rotary_emb,
             combined_img_seq_lengths,
         )
 
@@ -379,7 +420,7 @@ def _split_instruction_image_features(
 
 
 class BooguImageSelfAttention(nn.Module):
-    """GQA self-attention with QK RMSNorm and complex RoPE.
+    """GQA self-attention with QK RMSNorm and real-valued RoPE.
 
     Replaces the upstream diffusers `Attention` + Boogu attention processors;
     the vLLM-Omni `Attention` layer picks the kernel backend and handles
@@ -413,7 +454,7 @@ class BooguImageSelfAttention(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None,
-        rotary_emb: torch.Tensor | None,
+        rotary_emb: RotaryEmbedding | None,
     ) -> torch.Tensor:
         dtype = hidden_states.dtype
 
@@ -490,7 +531,7 @@ class BooguImageJointAttention(nn.Module):
         img_hidden_states: torch.Tensor,
         instruct_hidden_states: torch.Tensor,
         joint_attention_mask: torch.Tensor | None,
-        rotary_emb: torch.Tensor | None,
+        rotary_emb: RotaryEmbedding | None,
         encoder_seq_lengths: list[int],
         seq_lengths: list[int],
     ) -> torch.Tensor:
@@ -579,7 +620,7 @@ class BooguImageTransformerBlock(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None,
-        image_rotary_emb: torch.Tensor | None,
+        image_rotary_emb: RotaryEmbedding | None,
         temb: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if self.modulation:
@@ -687,8 +728,8 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module):
         instruct_hidden_states: torch.Tensor,  # [B, L_instruct, D]
         img_attention_mask: torch.Tensor | None,  # [B, L_img]
         joint_attention_mask: torch.Tensor | None,  # [B, L_total]
-        image_rotary_emb: torch.Tensor | None,  # [B, L_img, head_dim // 2]
-        rotary_emb: torch.Tensor | None,  # [B, L_total, head_dim // 2]
+        image_rotary_emb: RotaryEmbedding | None,  # Each tensor is [B, L_img, head_dim].
+        rotary_emb: RotaryEmbedding | None,  # Each tensor is [B, L_total, head_dim].
         temb: torch.Tensor | None = None,
         encoder_seq_lengths: list[int] | None = None,
         seq_lengths: list[int] | None = None,
@@ -1103,15 +1144,15 @@ class BooguImageTransformer2DModel(nn.Module):
 
     def img_patch_embed_and_refine(
         self,
-        hidden_states,
-        ref_image_hidden_states,
-        padded_img_mask,
-        padded_ref_img_mask,
-        noise_rotary_emb,
-        ref_img_rotary_emb,
-        l_effective_ref_img_len,
-        l_effective_img_len,
-        temb,
+        hidden_states: torch.Tensor,
+        ref_image_hidden_states: torch.Tensor,
+        padded_img_mask: torch.Tensor,
+        padded_ref_img_mask: torch.Tensor,
+        noise_rotary_emb: RotaryEmbedding,
+        ref_img_rotary_emb: RotaryEmbedding,
+        l_effective_ref_img_len: list[list[int]],
+        l_effective_img_len: list[int],
+        temb: torch.Tensor,
     ):
         """Embed image patches and run the refiner blocks.
 
@@ -1148,8 +1189,19 @@ class BooguImageTransformer2DModel(nn.Module):
             batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros(
                 num_ref_images, max_ref_img_len, self.hidden_size
             )
-            batch_ref_img_rotary_emb = hidden_states.new_zeros(
-                num_ref_images, max_ref_img_len, ref_img_rotary_emb.shape[-1], dtype=ref_img_rotary_emb.dtype
+            batch_ref_img_rotary_emb: RotaryEmbedding = (
+                hidden_states.new_zeros(
+                    num_ref_images,
+                    max_ref_img_len,
+                    ref_img_rotary_emb[0].shape[-1],
+                    dtype=ref_img_rotary_emb[0].dtype,
+                ),
+                hidden_states.new_zeros(
+                    num_ref_images,
+                    max_ref_img_len,
+                    ref_img_rotary_emb[1].shape[-1],
+                    dtype=ref_img_rotary_emb[1].dtype,
+                ),
             )
             batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype)
 
@@ -1162,7 +1214,8 @@ class BooguImageTransformer2DModel(nn.Module):
                     batch_ref_image_hidden_states[idx, :ref_img_len] = ref_image_hidden_states[
                         i, shift : shift + ref_img_len
                     ]
-                    batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[i, shift : shift + ref_img_len]
+                    for batch_freqs, ref_freqs in zip(batch_ref_img_rotary_emb, ref_img_rotary_emb):
+                        batch_freqs[idx, :ref_img_len] = ref_freqs[i, shift : shift + ref_img_len]
                     batch_temb[idx] = temb[i]
                     shift += ref_img_len
                     idx += 1
@@ -1195,7 +1248,7 @@ class BooguImageTransformer2DModel(nn.Module):
         hidden_states: torch.Tensor | list[torch.Tensor],
         timestep: torch.Tensor,
         instruction_hidden_states: torch.Tensor,
-        freqs_cis: torch.Tensor,
+        freqs_real: RotaryFrequencyTables,
         instruction_attention_mask: torch.Tensor,
         ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
     ) -> torch.Tensor:
@@ -1243,7 +1296,7 @@ class BooguImageTransformer2DModel(nn.Module):
             combined_img_rotary_emb,
             combined_img_seq_lengths,
         ) = self.rope_embedder(
-            freqs_cis,
+            freqs_real,
             instruction_attention_mask,
             l_effective_ref_img_len,
             l_effective_img_len,

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -544,7 +544,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
         return height, width, ori_height, ori_width
 
     def predict(
-        self, t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, ref_image_hidden_states=None
+        self, t, latents, instruction_embeds, freqs_real, instruction_attention_mask, ref_image_hidden_states=None
     ):
         """One transformer velocity prediction (upstream ``predict``).
 
@@ -557,7 +557,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
             latents,
             timestep,
             instruction_embeds,
-            freqs_cis,
+            freqs_real,
             instruction_attention_mask,
             ref_image_hidden_states=ref_image_hidden_states,
         )
@@ -710,7 +710,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
             generator,
         )
 
-        freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(
+        freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
             self.transformer.axes_dim_rope,
             self.transformer.axes_lens,
             theta=10000,
@@ -733,7 +733,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                 image_gs = image_guidance_scale if in_cfg_range else 1.0
 
                 model_pred = self.predict(
-                    t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, ref_latents
+                    t, latents, instruction_embeds, freqs_real, instruction_attention_mask, ref_latents
                 )
 
                 if task_type == "ti2i" and text_gs > 1.0 and image_gs > 1.0:
@@ -742,7 +742,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                         t,
                         latents,
                         negative_instruction_embeds,
-                        freqs_cis,
+                        freqs_real,
                         negative_instruction_attention_mask,
                         ref_latents,
                     )
@@ -750,7 +750,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                         t,
                         latents,
                         negative_instruction_embeds,
-                        freqs_cis,
+                        freqs_real,
                         negative_instruction_attention_mask,
                         None,
                     )
@@ -763,7 +763,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                         t,
                         latents,
                         negative_instruction_embeds,
-                        freqs_cis,
+                        freqs_real,
                         negative_instruction_attention_mask,
                         ref_latents,
                     )
@@ -771,7 +771,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                 elif task_type == "ti2i" and image_gs > 1.0:
                     # Image-only ti2i guidance: drop the reference in the uncond pred.
                     model_pred_drop_image = self.predict(
-                        t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, None
+                        t, latents, instruction_embeds, freqs_real, instruction_attention_mask, None
                     )
                     model_pred = model_pred + (image_gs - 1) * (model_pred - model_pred_drop_image)
                 elif text_gs > 1.0:
@@ -780,7 +780,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                         t,
                         latents,
                         negative_instruction_embeds,
-                        freqs_cis,
+                        freqs_real,
                         negative_instruction_attention_mask,
                         None,
                     )