# Copyright 2025 The JoyImage Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect import math import torch import torch.nn as nn import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin from ..normalization import FP32LayerNorm logger = logging.get_logger(__name__) # pylint: disable=invalid-name def _apply_rotary_emb_batched( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: """RoPE for batched [B, S, D] freqs.""" cos, sin = freqs_cis[0].to(xq.device), freqs_cis[1].to(xq.device) # batched: [B, S, D] -> [B, S, 1, D] cos = cos.unsqueeze(2) sin = sin.unsqueeze(2) def _rotate_half(x): x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) return torch.stack([-x_imag, x_real], dim=-1).flatten(3) xq_out = (xq.float() * cos + _rotate_half(xq) * sin).type_as(xq) xk_out = (xk.float() * cos + _rotate_half(xk) * sin).type_as(xk) return xq_out, xk_out # Copied from diffusers.models.transformers.transformer_joyimage.JoyImageModulate with JoyImage->JoyImageEditPlus class JoyImageEditPlusModulate(nn.Module): """Wan-style learnable modulation table. Produces `factor` modulation vectors by adding the conditioning signal to a learnable parameter table. """ def __init__(self, hidden_size: int, factor: int, dtype=None, device=None): super().__init__() self.factor = factor self.modulate_table = nn.Parameter( torch.zeros(1, factor, hidden_size, dtype=dtype, device=device) / hidden_size**0.5, requires_grad=True, ) def forward(self, x: torch.Tensor) -> list[torch.Tensor]: if x.ndim != 3: x = x.unsqueeze(1) return [o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1)] class JoyImageEditPlusAttnProcessor: """Attention processor that supports batched RoPE embeddings for edit-plus multi-image input.""" _attention_backend = None _parallel_config = None def __call__( self, attn: "JoyImageEditPlusAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if encoder_hidden_states is None: raise ValueError("JoyImageEditPlusAttnProcessor requires encoder_hidden_states") heads = attn.heads img_qkv = attn.img_attn_qkv(hidden_states) img_query, img_key, img_value = img_qkv.chunk(3, dim=-1) txt_qkv = attn.txt_attn_qkv(encoder_hidden_states) txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1) img_query = img_query.unflatten(-1, (heads, -1)) img_key = img_key.unflatten(-1, (heads, -1)) img_value = img_value.unflatten(-1, (heads, -1)) txt_query = txt_query.unflatten(-1, (heads, -1)) txt_key = txt_key.unflatten(-1, (heads, -1)) txt_value = txt_value.unflatten(-1, (heads, -1)) img_query = attn.img_attn_q_norm(img_query) img_key = attn.img_attn_k_norm(img_key) txt_query = attn.txt_attn_q_norm(txt_query) txt_key = attn.txt_attn_k_norm(txt_key) if image_rotary_emb is not None: img_query, img_key = _apply_rotary_emb_batched(img_query, img_key, image_rotary_emb) joint_query = torch.cat([img_query, txt_query], dim=1) joint_key = torch.cat([img_key, txt_key], dim=1) joint_value = torch.cat([img_value, txt_value], dim=1) joint_hidden_states = dispatch_attention_fn( joint_query, joint_key, joint_value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, backend=self._attention_backend, parallel_config=self._parallel_config, ) joint_hidden_states = joint_hidden_states.flatten(2, 3) joint_hidden_states = joint_hidden_states.to(joint_query.dtype) img_attn_output = joint_hidden_states[:, : hidden_states.shape[1], :] txt_attn_output = joint_hidden_states[:, hidden_states.shape[1] :, :] img_attn_output = attn.img_attn_proj(img_attn_output) txt_attn_output = attn.txt_attn_proj(txt_attn_output) return img_attn_output, txt_attn_output class JoyImageEditPlusAttention(nn.Module, AttentionModuleMixin): """Joint attention module for JoyImage Edit Plus double-stream blocks.""" _default_processor_cls = JoyImageEditPlusAttnProcessor _available_processors = [JoyImageEditPlusAttnProcessor] _supports_qkv_fusion = False def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, eps: float = 1e-6, processor=None, ): super().__init__() self.heads = num_attention_heads self.head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim self.img_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True) self.img_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps) self.img_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps) self.img_attn_proj = nn.Linear(inner_dim, dim, bias=True) self.txt_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True) self.txt_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps) self.txt_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps) self.txt_attn_proj = nn.Linear(inner_dim, dim, bias=True) if processor is None: processor = self._default_processor_cls() self.set_processor(processor) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) kwargs = {} if "attention_mask" in attn_parameters: kwargs["attention_mask"] = attention_mask return self.processor(self, hidden_states, encoder_hidden_states, image_rotary_emb, **kwargs) class JoyImageEditPlusTransformerBlock(nn.Module): """Double-stream transformer block for JoyImage Edit Plus.""" def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_width_ratio: float = 4.0, eps: float = 1e-6, ): super().__init__() self.dim = dim self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim mlp_hidden_dim = int(dim * mlp_width_ratio) # image stream self.img_mod = JoyImageEditPlusModulate(dim, factor=6) self.img_norm1 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) self.img_norm2 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) self.img_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate") # text stream self.txt_mod = JoyImageEditPlusModulate(dim, factor=6) self.txt_norm1 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) self.txt_norm2 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps) self.txt_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate") # joint attention self.attn = JoyImageEditPlusAttention(dim, num_attention_heads, attention_head_dim, eps=eps) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # modulation ( img_mod1_shift, img_mod1_scale, img_mod1_gate, img_mod2_shift, img_mod2_scale, img_mod2_gate, ) = self.img_mod(temb) ( txt_mod1_shift, txt_mod1_scale, txt_mod1_gate, txt_mod2_shift, txt_mod2_scale, txt_mod2_gate, ) = self.txt_mod(temb) # --- attention --- img_normed = self.img_norm1(hidden_states) txt_normed = self.txt_norm1(encoder_hidden_states) img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1) txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1) img_attn, txt_attn = self.attn( hidden_states=img_modulated, encoder_hidden_states=txt_modulated, image_rotary_emb=image_rotary_emb, attention_mask=attention_mask, ) hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1) # --- FFN --- img_ffn_normed = self.img_norm2(hidden_states) txt_ffn_normed = self.txt_norm2(encoder_hidden_states) img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1) txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1) img_ffn_output = self.img_mlp(img_ffn_input) txt_ffn_output = self.txt_mlp(txt_ffn_input) hidden_states = hidden_states + img_ffn_output * img_mod2_gate.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + txt_ffn_output * txt_mod2_gate.unsqueeze(1) return hidden_states, encoder_hidden_states # Copied from diffusers.models.transformers.transformer_joyimage.JoyImageTimeTextImageEmbedding with JoyImage->JoyImageEditPlus class JoyImageEditPlusTimeTextImageEmbedding(nn.Module): def __init__( self, dim: int, time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, ): super().__init__() self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) self.act_fn = nn.SiLU() self.time_proj = nn.Linear(dim, time_proj_dim) self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, ): timestep = self.timesteps_proj(timestep) time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: timestep = timestep.to(time_embedder_dtype) temb = self.time_embedder(timestep).type_as(encoder_hidden_states) timestep_proj = self.time_proj(self.act_fn(temb)) encoder_hidden_states = self.text_embedder(encoder_hidden_states) return temb, timestep_proj, encoder_hidden_states class JoyImageEditPlusTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin): r""" JoyImage Edit Plus Transformer for multi-image editing. Uses a patchify+padding approach where each reference image and the target noise are independently patchified and concatenated into a flat patch sequence. Supports variable-resolution reference images. Input format: `[B, max_patches, C, pt, ph, pw]` (6D padded patches). Args: patch_size (`list`, defaults to `[1, 2, 2]`): Patch size for patchifying the latent input along `(t, h, w)` dimensions. in_channels (`int`, defaults to `16`): The number of channels in the input latent. out_channels (`int`, *optional*, defaults to `None`): The number of channels in the output. If not specified, it defaults to `in_channels`. hidden_size (`int`, defaults to `3072`): The dimensionality of the hidden representations. num_attention_heads (`int`, defaults to `24`): The number of attention heads. text_dim (`int`, defaults to `4096`): The dimensionality of the text encoder output. mlp_width_ratio (`float`, defaults to `4.0`): The ratio of MLP hidden dimension to `hidden_size`. num_layers (`int`, defaults to `20`): The number of double-stream transformer blocks. rope_dim_list (`list[int]`, defaults to `[16, 56, 56]`): The dimensions for 3D rotary positional embeddings along `(t, h, w)`. rope_type (`str`, defaults to `"rope"`): The type of rotary positional embedding. theta (`int`, defaults to `256`): The base frequency for rotary embeddings. """ _skip_layerwise_casting_patterns = ["img_in", "condition_embedder", "norm"] _no_split_modules = ["JoyImageEditPlusTransformerBlock"] _supports_gradient_checkpointing = True _keep_in_fp32_modules = [ "time_embedder", "norm1", "norm2", "norm_out", ] _repeated_blocks = ["JoyImageEditPlusTransformerBlock"] @register_to_config def __init__( self, patch_size: list[int] = [1, 2, 2], in_channels: int = 16, out_channels: int | None = None, hidden_size: int = 3072, num_attention_heads: int = 24, text_dim: int = 4096, mlp_width_ratio: float = 4.0, num_layers: int = 20, rope_dim_list: list[int] = [16, 56, 56], rope_type: str = "rope", theta: int = 256, ): super().__init__() self.out_channels = out_channels or in_channels attention_head_dim = hidden_size // num_attention_heads if hidden_size % num_attention_heads != 0: raise ValueError( f"hidden_size ({hidden_size}) must be divisible by num_attention_heads ({num_attention_heads})" ) self.img_in = nn.Conv3d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size) self.condition_embedder = JoyImageEditPlusTimeTextImageEmbedding( dim=hidden_size, time_freq_dim=256, time_proj_dim=hidden_size * 6, text_embed_dim=text_dim, ) self.double_blocks = nn.ModuleList( [ JoyImageEditPlusTransformerBlock( dim=hidden_size, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, mlp_width_ratio=mlp_width_ratio, ) for _ in range(num_layers) ] ) self.norm_out = FP32LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.proj_out = nn.Linear(hidden_size, self.out_channels * math.prod(patch_size)) self.gradient_checkpointing = False # Set batched-RoPE-aware attention processor on all blocks for block in self.double_blocks: block.attn.set_processor(JoyImageEditPlusAttnProcessor()) def _get_rotary_pos_embed_for_range( self, start: tuple[int, int, int], stop: tuple[int, int, int], ) -> tuple[torch.Tensor, torch.Tensor]: """Generate 3D RoPE for a spatial range [start, stop).""" head_dim = self.config.hidden_size // self.config.num_attention_heads rope_dim_list = self.config.rope_dim_list if rope_dim_list is None: rope_dim_list = [head_dim // 3] * 3 grids = [] for i in range(3): grids.append(torch.arange(start[i], stop[i], dtype=torch.float32)) mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0) cos_parts, sin_parts = [], [] for i, dim in enumerate(rope_dim_list): pos = mesh[i].reshape(-1) freqs = 1.0 / (self.config.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)) angles = torch.outer(pos, freqs) cos_parts.append(angles.cos().repeat_interleave(2, dim=1)) sin_parts.append(angles.sin().repeat_interleave(2, dim=1)) return torch.cat(cos_parts, dim=1), torch.cat(sin_parts, dim=1) def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor | None = None, shape_list: list[list[tuple[int, int, int]]] = None, return_dict: bool = True, ) -> torch.Tensor | tuple: """ Args: hidden_states: [B, max_patches, C, pt, ph, pw] - patchified latent input. timestep: [B] - diffusion timestep. encoder_hidden_states: [B, L, D] - text encoder outputs. encoder_hidden_states_mask: [B, L] - attention mask for text tokens. shape_list: Per-sample list of (t, h, w) tuples for each component (target + references). return_dict: Whether to return a dict or tuple. Returns: If `return_dict` is True, an [`~models.modeling_outputs.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ batch_size, max_num_patches, channels, pt, ph, pw = hidden_states.shape device = hidden_states.device # 1. Condition embeddings _, vec, txt = self.condition_embedder(timestep, encoder_hidden_states) vec = vec.unflatten(1, (6, -1)) # 2. Patchify via Conv3d: flatten (B, N) -> apply conv -> reshape back x = hidden_states.reshape(batch_size * max_num_patches, channels, pt, ph, pw) x = self.img_in(x) # (B*N, D, 1, 1, 1) img = x.reshape(batch_size, max_num_patches, -1) # 3. Build per-component RoPE with temporal offsets sample_cos_list, sample_sin_list = [], [] for i in range(batch_size): s_cos_parts, s_sin_parts = [], [] current_t_offset = 0 for thw in shape_list[i]: t, h, w = thw start = (current_t_offset, 0, 0) stop = (current_t_offset + t, h, w) cos_emb, sin_emb = self._get_rotary_pos_embed_for_range(start, stop) s_cos_parts.append(cos_emb) s_sin_parts.append(sin_emb) current_t_offset += t s_cos = torch.cat(s_cos_parts, dim=0).to(device) s_sin = torch.cat(s_sin_parts, dim=0).to(device) actual_len = s_cos.shape[0] pad_len = max_num_patches - actual_len if pad_len > 0: s_cos = F.pad(s_cos, (0, 0, 0, pad_len), value=1.0) s_sin = F.pad(s_sin, (0, 0, 0, pad_len), value=0.0) sample_cos_list.append(s_cos) sample_sin_list.append(s_sin) vis_freqs = (torch.stack(sample_cos_list), torch.stack(sample_sin_list)) # 4. Build attention mask: [B, 1, 1, img_seq + txt_seq] attention_mask = None if encoder_hidden_states_mask is not None: img_mask = torch.zeros(batch_size, max_num_patches, device=device, dtype=encoder_hidden_states_mask.dtype) for i in range(batch_size): actual_len = sum(t * h * w for t, h, w in shape_list[i]) img_mask[i, :actual_len] = 1.0 full_mask = torch.cat([img_mask, encoder_hidden_states_mask], dim=1) attention_mask = full_mask.unsqueeze(1).unsqueeze(1).bool() # 5. Run double blocks for block in self.double_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: img, txt = self._gradient_checkpointing_func(block, img, txt, vec, vis_freqs, attention_mask) else: img, txt = block( hidden_states=img, encoder_hidden_states=txt, temb=vec, image_rotary_emb=vis_freqs, attention_mask=attention_mask, ) # 6. Output projection + reshape to 6D patches img = self.proj_out(self.norm_out(img)) img = img.reshape(batch_size, max_num_patches, pt, ph, pw, self.out_channels).permute( 0, 1, 5, 2, 3, 4 ) # -> [B, N, C, pt, ph, pw] if not return_dict: return (img,) return Transformer2DModelOutput(sample=img)