diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 08ba4af09..19b7ef04c 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/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, ) diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index 999fdbb41..84bac1b2b 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/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) diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index e7489c868..7016b4f69 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/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, diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index a3b56051d..02b333f4b 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/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, )