diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 2d282c2f9..96e4ad0cc 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -6,7 +6,7 @@ from typing import Any import torch import torch.nn as nn -from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle from vllm.logger import init_logger @@ -17,6 +17,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.utils.platform_utils import is_npu logger = init_logger(__name__) @@ -97,6 +98,7 @@ class LongCatImageAttention(nn.Module): self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) + self.rope = RotaryEmbedding(is_neox_style=False) self.attn = Attention( num_heads=heads, head_size=self.head_dim, @@ -138,8 +140,11 @@ class LongCatImageAttention(nn.Module): value = torch.cat([encoder_value, value], dim=1) if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) - key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + cos, sin = image_rotary_emb # [S, D/2] + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) hidden_states = self.attn( query, @@ -263,16 +268,15 @@ class LongCatImagePosEmbed(nn.Module): is_npu = ids.device.type == "npu" freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 for i in range(n_axes): - cos, sin = get_1d_rotary_pos_embed( + freqs_cis = get_1d_rotary_pos_embed( self.axes_dim[i], pos[:, i], theta=self.theta, - repeat_interleave_real=True, - use_real=True, + use_real=False, freqs_dtype=freqs_dtype, ) - cos_out.append(cos) - sin_out.append(sin) + cos_out.append(freqs_cis.real) + sin_out.append(freqs_cis.imag) freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) return freqs_cos, freqs_sin diff --git a/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py index ea3916b7a..8a4629e98 100644 --- a/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py +++ b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py @@ -21,7 +21,7 @@ from typing import Any import torch import torch.nn as nn from diffusers.models.attention import FeedForward -from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle from diffusers.utils import is_torch_npu_available @@ -32,6 +32,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -96,6 +97,7 @@ class OvisImageAttention(nn.Module): self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) + self.rope = RotaryEmbedding(is_neox_style=False) self.attn = Attention( num_heads=heads, head_size=self.head_dim, @@ -137,8 +139,11 @@ class OvisImageAttention(nn.Module): value = torch.cat([encoder_value, value], dim=1) if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) - key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + cos, sin = image_rotary_emb # [S, D/2] + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) hidden_states = self.attn( query, @@ -318,16 +323,15 @@ class OvisImagePosEmbed(nn.Module): is_npu = ids.device.type == "npu" freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 for i in range(n_axes): - cos, sin = get_1d_rotary_pos_embed( + freqs_cis = get_1d_rotary_pos_embed( self.axes_dim[i], pos[:, i], theta=self.theta, - repeat_interleave_real=True, - use_real=True, + use_real=False, freqs_dtype=freqs_dtype, ) - cos_out.append(cos) - sin_out.append(sin) + cos_out.append(freqs_cis.real) + sin_out.append(freqs_cis.imag) freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) return freqs_cos, freqs_sin