Add support for Boogu Image and Boogu Image Edit
This commit is contained in:
Jaret Burkett
@@ -0,0 +1,491 @@
|
||||
# Vendored from the Boogu-Image repository (boogu/models/attention_processor.py).
|
||||
# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0.
|
||||
#
|
||||
# Attention here defaults to torch's ``scaled_dot_product_attention`` (the
|
||||
# "native" backend) so the model has NO hard dependency on flash-attn. Flash
|
||||
# Attention 2 is an OPTIONAL backend: each processor carries an
|
||||
# ``attention_backend`` flag (set in bulk via
|
||||
# ``BooguImageTransformer2DModel.set_attention_backend``) and only the "flash"
|
||||
# branch touches the ``flash_attn`` package, so importing it stays lazy/guarded.
|
||||
import math
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.attention_processor import Attention
|
||||
from einops import repeat
|
||||
|
||||
from .embeddings import apply_rotary_emb
|
||||
|
||||
try:
|
||||
from flash_attn import flash_attn_varlen_func
|
||||
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
|
||||
|
||||
_FLASH_ATTN_AVAILABLE = True
|
||||
except ImportError: # flash-attn is optional; "native" SDPA needs none of this.
|
||||
flash_attn_varlen_func = None
|
||||
index_first_axis = pad_input = unpad_input = None
|
||||
_FLASH_ATTN_AVAILABLE = False
|
||||
|
||||
# Supported attention backends. "native" -> SDPA, "flash" -> Flash Attention 2.
|
||||
ATTENTION_BACKENDS = ("native", "flash")
|
||||
|
||||
|
||||
def _get_unpad_data(mask_2d: torch.Tensor):
|
||||
"""Indices / cu_seqlens / max_seqlen from a 2D padding mask [B, L]."""
|
||||
seqlens_in_batch = mask_2d.sum(dim=-1, dtype=torch.int32)
|
||||
indices = torch.nonzero(mask_2d.flatten(), as_tuple=False).flatten()
|
||||
max_seqlen_in_batch = seqlens_in_batch.max().item()
|
||||
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
|
||||
return indices, cu_seqlens, max_seqlen_in_batch
|
||||
|
||||
|
||||
def _upad_input(query, key, value, attention_mask, query_length, num_heads):
|
||||
"""Unpad q/k/v for ``flash_attn_varlen_func`` given a [B, L] padding mask."""
|
||||
indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
|
||||
batch_size, kv_seq_len, num_key_value_heads, head_dim = key.shape
|
||||
|
||||
key = index_first_axis(
|
||||
key.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
|
||||
)
|
||||
value = index_first_axis(
|
||||
value.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
|
||||
)
|
||||
|
||||
if query_length == kv_seq_len:
|
||||
query = index_first_axis(
|
||||
query.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
|
||||
)
|
||||
cu_seqlens_q = cu_seqlens_k
|
||||
max_seqlen_in_batch_q = max_seqlen_in_batch_k
|
||||
indices_q = indices_k
|
||||
elif query_length == 1:
|
||||
max_seqlen_in_batch_q = 1
|
||||
cu_seqlens_q = torch.arange(
|
||||
batch_size + 1, dtype=torch.int32, device=query.device
|
||||
)
|
||||
indices_q = cu_seqlens_q[:-1]
|
||||
query = query.squeeze(1)
|
||||
else:
|
||||
q_mask = attention_mask[:, -query_length:]
|
||||
query, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
|
||||
query, q_mask
|
||||
)
|
||||
|
||||
return (
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
indices_q,
|
||||
(cu_seqlens_q, cu_seqlens_k),
|
||||
(max_seqlen_in_batch_q, max_seqlen_in_batch_k),
|
||||
)
|
||||
|
||||
|
||||
def _flash_varlen_attention(query, key, value, attention_mask, attn, softmax_scale):
|
||||
"""Run flash-attn varlen over a [B, L, heads, head_dim] q/k/v with a 2D mask.
|
||||
|
||||
Returns the attention output flattened back to [B, L, heads * head_dim].
|
||||
"""
|
||||
batch_size, sequence_length = query.shape[0], query.shape[1]
|
||||
kv_heads = key.shape[2]
|
||||
|
||||
mask_2d = attention_mask.bool() if attention_mask is not None else None
|
||||
(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
indices_q,
|
||||
(cu_seqlens_q, cu_seqlens_k),
|
||||
(max_seqlen_q, max_seqlen_k),
|
||||
) = _upad_input(query, key, value, mask_2d, sequence_length, attn.heads)
|
||||
|
||||
if kv_heads < attn.heads:
|
||||
key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
|
||||
value_states = repeat(
|
||||
value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads
|
||||
)
|
||||
|
||||
attn_output_unpad = flash_attn_varlen_func(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
dropout_p=0.0,
|
||||
causal=False,
|
||||
softmax_scale=softmax_scale,
|
||||
)
|
||||
hidden_states = pad_input(attn_output_unpad, indices_q, batch_size, sequence_length)
|
||||
return hidden_states.flatten(-2)
|
||||
|
||||
|
||||
class BooguImageDoubleStreamSelfAttnProcessor(nn.Module):
|
||||
"""
|
||||
Double-stream self-attention processor.
|
||||
|
||||
Instruction and image features each get their own q/k/v projections; the two
|
||||
streams are concatenated (instruction first), attended jointly, then split
|
||||
back and projected with separate output heads. Uses torch SDPA by default;
|
||||
set ``attention_backend = "flash"`` for Flash Attention 2.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
num_attention_heads: int,
|
||||
num_kv_heads: int,
|
||||
qkv_bias: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError(
|
||||
"BooguImageDoubleStreamSelfAttnProcessor requires PyTorch 2.0+."
|
||||
)
|
||||
|
||||
self.head_dim = head_dim
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.attention_backend = "native"
|
||||
|
||||
query_dim = head_dim * num_attention_heads
|
||||
kv_dim = head_dim * num_kv_heads
|
||||
|
||||
self.img_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias)
|
||||
self.img_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
|
||||
self.img_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
|
||||
|
||||
self.instruct_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias)
|
||||
self.instruct_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
|
||||
self.instruct_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
|
||||
|
||||
self.instruct_out = nn.Linear(query_dim, query_dim, bias=qkv_bias)
|
||||
self.img_out = nn.Linear(query_dim, query_dim, bias=qkv_bias)
|
||||
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self) -> None:
|
||||
nn.init.xavier_uniform_(self.img_to_q.weight)
|
||||
nn.init.xavier_uniform_(self.img_to_k.weight)
|
||||
nn.init.xavier_uniform_(self.img_to_v.weight)
|
||||
nn.init.xavier_uniform_(self.instruct_to_q.weight)
|
||||
nn.init.xavier_uniform_(self.instruct_to_k.weight)
|
||||
nn.init.xavier_uniform_(self.instruct_to_v.weight)
|
||||
nn.init.xavier_uniform_(self.instruct_out.weight)
|
||||
nn.init.xavier_uniform_(self.img_out.weight)
|
||||
|
||||
if self.img_to_q.bias is not None:
|
||||
nn.init.zeros_(self.img_to_q.bias)
|
||||
nn.init.zeros_(self.img_to_k.bias)
|
||||
nn.init.zeros_(self.img_to_v.bias)
|
||||
nn.init.zeros_(self.instruct_to_q.bias)
|
||||
nn.init.zeros_(self.instruct_to_k.bias)
|
||||
nn.init.zeros_(self.instruct_to_v.bias)
|
||||
nn.init.zeros_(self.instruct_out.bias)
|
||||
nn.init.zeros_(self.img_out.bias)
|
||||
|
||||
def _concat_instruction_image_features(
|
||||
self,
|
||||
img_hidden_states_list: List[torch.Tensor],
|
||||
instruct_hidden_states_list: List[torch.Tensor],
|
||||
encoder_seq_lengths: List[int],
|
||||
seq_lengths: List[int],
|
||||
) -> List[torch.Tensor]:
|
||||
"""Concatenate instruction then image features into one joint sequence."""
|
||||
batch_size = img_hidden_states_list[0].shape[0]
|
||||
max_seq_len = max(seq_lengths)
|
||||
|
||||
concatenated_list = []
|
||||
for img_tensor, instruct_tensor in zip(
|
||||
img_hidden_states_list, instruct_hidden_states_list
|
||||
):
|
||||
device = img_tensor.device
|
||||
if instruct_tensor.device != device:
|
||||
instruct_tensor = instruct_tensor.to(device)
|
||||
|
||||
feature_dim = img_tensor.shape[-1]
|
||||
concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim)
|
||||
|
||||
for i, (encoder_seq_len, seq_len) in enumerate(
|
||||
zip(encoder_seq_lengths, seq_lengths)
|
||||
):
|
||||
concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len]
|
||||
concatenated[i, encoder_seq_len:seq_len] = img_tensor[
|
||||
i, : seq_len - encoder_seq_len
|
||||
]
|
||||
|
||||
concatenated_list.append(concatenated)
|
||||
|
||||
return concatenated_list
|
||||
|
||||
def _split_instruction_image_features(
|
||||
self,
|
||||
hidden_states_list: List[torch.Tensor],
|
||||
encoder_seq_lengths: List[int],
|
||||
seq_lengths: List[int],
|
||||
) -> List[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Inverse of ``_concat_instruction_image_features``."""
|
||||
result_list = []
|
||||
for hidden_states in hidden_states_list:
|
||||
batch_size = hidden_states.shape[0]
|
||||
feature_dim = hidden_states.shape[-1]
|
||||
|
||||
max_instruct_len = max(encoder_seq_lengths)
|
||||
max_img_len = max(
|
||||
seq_len - encoder_seq_len
|
||||
for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths)
|
||||
)
|
||||
|
||||
instruct_hidden_states = hidden_states.new_zeros(
|
||||
batch_size, max_instruct_len, feature_dim
|
||||
)
|
||||
img_hidden_states = hidden_states.new_zeros(
|
||||
batch_size, max_img_len, feature_dim
|
||||
)
|
||||
|
||||
for i, (encoder_seq_len, seq_len) in enumerate(
|
||||
zip(encoder_seq_lengths, seq_lengths)
|
||||
):
|
||||
img_len = seq_len - encoder_seq_len
|
||||
instruct_hidden_states[i, :encoder_seq_len] = hidden_states[
|
||||
i, :encoder_seq_len
|
||||
]
|
||||
img_hidden_states[i, :img_len] = hidden_states[
|
||||
i, encoder_seq_len:seq_len
|
||||
]
|
||||
|
||||
result_list.append((instruct_hidden_states, img_hidden_states))
|
||||
|
||||
return result_list
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
img_hidden_states: torch.Tensor,
|
||||
instruct_hidden_states: torch.Tensor,
|
||||
joint_attention_mask: Optional[torch.Tensor] = None,
|
||||
rotary_emb: Optional[torch.Tensor] = None,
|
||||
encoder_seq_lengths: List[int] = None,
|
||||
seq_lengths: List[int] = None,
|
||||
base_sequence_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
batch_size = img_hidden_states.shape[0]
|
||||
|
||||
img_query = self.img_to_q(img_hidden_states)
|
||||
img_key = self.img_to_k(img_hidden_states)
|
||||
img_value = self.img_to_v(img_hidden_states)
|
||||
|
||||
instruct_query = self.instruct_to_q(instruct_hidden_states)
|
||||
instruct_key = self.instruct_to_k(instruct_hidden_states)
|
||||
instruct_value = self.instruct_to_v(instruct_hidden_states)
|
||||
|
||||
img_list = [img_query, img_key, img_value]
|
||||
instruct_list = [instruct_query, instruct_key, instruct_value]
|
||||
concatenated_list = self._concat_instruction_image_features(
|
||||
img_list, instruct_list, encoder_seq_lengths, seq_lengths
|
||||
)
|
||||
query, key, value = concatenated_list
|
||||
|
||||
sequence_length = max(seq_lengths)
|
||||
|
||||
query_dim = query.shape[-1]
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = query_dim // attn.heads
|
||||
dtype = query.dtype
|
||||
|
||||
kv_heads = inner_dim // head_dim
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim)
|
||||
key = key.view(batch_size, -1, kv_heads, head_dim)
|
||||
value = value.view(batch_size, -1, kv_heads, head_dim)
|
||||
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
|
||||
if rotary_emb is not None:
|
||||
query = apply_rotary_emb(query, rotary_emb, use_real=False)
|
||||
key = apply_rotary_emb(key, rotary_emb, use_real=False)
|
||||
|
||||
query, key = query.to(dtype), key.to(dtype)
|
||||
|
||||
if base_sequence_length is not None:
|
||||
softmax_scale = (
|
||||
math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
|
||||
)
|
||||
else:
|
||||
softmax_scale = attn.scale
|
||||
|
||||
if self.attention_backend == "flash":
|
||||
# q/k/v are [B, L, heads, head_dim]; the joint padding mask is 2D.
|
||||
hidden_states = _flash_varlen_attention(
|
||||
query, key, value, joint_attention_mask, attn, softmax_scale
|
||||
)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
else:
|
||||
if joint_attention_mask is not None:
|
||||
joint_attention_mask = joint_attention_mask.bool()
|
||||
if joint_attention_mask.dim() == 2:
|
||||
joint_attention_mask = joint_attention_mask.view(
|
||||
batch_size, 1, 1, -1
|
||||
)
|
||||
elif joint_attention_mask.dim() == 3:
|
||||
joint_attention_mask = joint_attention_mask.unsqueeze(1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported joint_attention_mask shape: {joint_attention_mask.shape}"
|
||||
)
|
||||
|
||||
q = query.transpose(1, 2)
|
||||
k = key.transpose(1, 2)
|
||||
v = value.transpose(1, 2)
|
||||
|
||||
# explicitly repeat key/value to avoid the slow MATH SDPA backend that
|
||||
# enable_gqa triggers on some torch builds
|
||||
k = k.repeat_interleave(q.size(-3) // k.size(-3), -3)
|
||||
v = v.repeat_interleave(q.size(-3) // v.size(-3), -3)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=joint_attention_mask, scale=softmax_scale
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(
|
||||
batch_size, -1, attn.heads * head_dim
|
||||
)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
split_results = self._split_instruction_image_features(
|
||||
[hidden_states], encoder_seq_lengths, seq_lengths
|
||||
)
|
||||
instruct_hidden_states, img_hidden_states = split_results[0]
|
||||
|
||||
instruct_projected = self.instruct_out(instruct_hidden_states)
|
||||
img_projected = self.img_out(img_hidden_states)
|
||||
|
||||
merged_list = self._concat_instruction_image_features(
|
||||
[img_projected], [instruct_projected], encoder_seq_lengths, seq_lengths
|
||||
)
|
||||
hidden_states = merged_list[0]
|
||||
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class BooguImageAttnProcessor:
|
||||
"""
|
||||
Single-stream self-attention processor with RoPE + QK norm.
|
||||
|
||||
Uses torch SDPA by default; set ``attention_backend = "flash"`` for Flash
|
||||
Attention 2 (requires the ``flash_attn`` package).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("BooguImageAttnProcessor requires PyTorch 2.0+.")
|
||||
self.attention_backend = "native"
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
base_sequence_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
batch_size, sequence_length, _ = hidden_states.shape
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
query_dim = query.shape[-1]
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = query_dim // attn.heads
|
||||
dtype = query.dtype
|
||||
|
||||
kv_heads = inner_dim // head_dim
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim)
|
||||
key = key.view(batch_size, -1, kv_heads, head_dim)
|
||||
value = value.view(batch_size, -1, kv_heads, head_dim)
|
||||
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
query = apply_rotary_emb(query, image_rotary_emb, use_real=False)
|
||||
key = apply_rotary_emb(key, image_rotary_emb, use_real=False)
|
||||
|
||||
query, key = query.to(dtype), key.to(dtype)
|
||||
|
||||
if base_sequence_length is not None:
|
||||
softmax_scale = (
|
||||
math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
|
||||
)
|
||||
else:
|
||||
softmax_scale = attn.scale
|
||||
|
||||
if self.attention_backend == "flash" and (
|
||||
attention_mask is None or attention_mask.dim() == 2
|
||||
):
|
||||
mask = (
|
||||
attention_mask
|
||||
if attention_mask is not None
|
||||
else query.new_ones(batch_size, sequence_length, dtype=torch.bool)
|
||||
)
|
||||
hidden_states = _flash_varlen_attention(
|
||||
query, key, value, mask, attn, softmax_scale
|
||||
)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask.bool()
|
||||
if attention_mask.dim() == 2:
|
||||
attention_mask = attention_mask.view(batch_size, 1, 1, -1)
|
||||
elif attention_mask.dim() == 3:
|
||||
B, L, _ = attention_mask.shape
|
||||
diag_valid = torch.diagonal(attention_mask, dim1=-2, dim2=-1)
|
||||
lengths = diag_valid.sum(dim=-1)
|
||||
arange_L = torch.arange(L, device=attention_mask.device)
|
||||
q_valid = arange_L.unsqueeze(0) < lengths.unsqueeze(1)
|
||||
k_valid = q_valid
|
||||
causal = torch.tril(
|
||||
torch.ones(L, L, dtype=torch.bool, device=attention_mask.device)
|
||||
)
|
||||
combined = causal & q_valid.unsqueeze(-1) & k_valid.unsqueeze(-2)
|
||||
attention_mask = combined.unsqueeze(1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported attention_mask shape: {attention_mask.shape}"
|
||||
)
|
||||
|
||||
query = query.transpose(1, 2)
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
|
||||
key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
|
||||
value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, scale=softmax_scale
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(
|
||||
batch_size, -1, attn.heads * head_dim
|
||||
)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,164 @@
|
||||
# Vendored from the Boogu-Image repository (boogu/models/transformers/block_lumina2.py).
|
||||
# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0.
|
||||
#
|
||||
# The optional triton RMSNorm and flash-attn SwiGLU fast paths are dropped here;
|
||||
# we always use torch.nn.RMSNorm and a plain SwiGLU so the model runs anywhere.
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.embeddings import Timesteps
|
||||
from torch.nn import RMSNorm
|
||||
|
||||
from .embeddings import TimestepEmbedding
|
||||
|
||||
|
||||
def swiglu(x, y):
|
||||
return F.silu(x.float(), inplace=False).to(x.dtype) * y
|
||||
|
||||
|
||||
class LuminaRMSNormZero(nn.Module):
|
||||
"""Adaptive RMS normalization with a zero-initialized modulation projection."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
norm_eps: float,
|
||||
norm_elementwise_affine: bool,
|
||||
):
|
||||
super().__init__()
|
||||
self.silu = nn.SiLU()
|
||||
self.linear = nn.Linear(
|
||||
min(embedding_dim, 1024),
|
||||
4 * embedding_dim,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
self.norm = RMSNorm(embedding_dim, eps=norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
emb: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
emb = self.linear(self.silu(emb))
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
|
||||
x = self.norm(x) * (1 + scale_msa[:, None])
|
||||
return x, gate_msa, scale_mlp, gate_mlp
|
||||
|
||||
|
||||
class LuminaLayerNormContinuous(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
conditioning_embedding_dim: int,
|
||||
elementwise_affine=True,
|
||||
eps=1e-5,
|
||||
bias=True,
|
||||
norm_type="layer_norm",
|
||||
out_dim: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# AdaLN
|
||||
self.silu = nn.SiLU()
|
||||
self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
|
||||
|
||||
if norm_type == "layer_norm":
|
||||
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
|
||||
elif norm_type == "rms_norm":
|
||||
self.norm = RMSNorm(
|
||||
embedding_dim, eps=eps, elementwise_affine=elementwise_affine
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown norm_type {norm_type}")
|
||||
|
||||
self.linear_2 = None
|
||||
if out_dim is not None:
|
||||
self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
conditioning_embedding: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
|
||||
scale = emb
|
||||
x = self.norm(x) * (1 + scale)[:, None, :]
|
||||
|
||||
if self.linear_2 is not None:
|
||||
x = self.linear_2(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LuminaFeedForward(nn.Module):
|
||||
"""A SwiGLU feed-forward layer with a multiple-of-256 inner dim."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
inner_dim: int,
|
||||
multiple_of: Optional[int] = 256,
|
||||
ffn_dim_multiplier: Optional[float] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.swiglu = swiglu
|
||||
|
||||
if ffn_dim_multiplier is not None:
|
||||
inner_dim = int(ffn_dim_multiplier * inner_dim)
|
||||
inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
|
||||
|
||||
self.linear_1 = nn.Linear(dim, inner_dim, bias=False)
|
||||
self.linear_2 = nn.Linear(inner_dim, dim, bias=False)
|
||||
self.linear_3 = nn.Linear(dim, inner_dim, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
h1, h2 = self.linear_1(x), self.linear_3(x)
|
||||
return self.linear_2(self.swiglu(h1, h2))
|
||||
|
||||
|
||||
class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int = 4096,
|
||||
instruction_feat_dim: int = 2048,
|
||||
frequency_embedding_size: int = 256,
|
||||
norm_eps: float = 1e-5,
|
||||
timestep_scale: float = 1.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.time_proj = Timesteps(
|
||||
num_channels=frequency_embedding_size,
|
||||
flip_sin_to_cos=True,
|
||||
downscale_freq_shift=0.0,
|
||||
scale=timestep_scale,
|
||||
)
|
||||
|
||||
self.timestep_embedder = TimestepEmbedding(
|
||||
in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
|
||||
)
|
||||
|
||||
self.caption_embedder = nn.Sequential(
|
||||
RMSNorm(instruction_feat_dim, eps=norm_eps),
|
||||
nn.Linear(instruction_feat_dim, hidden_size, bias=True),
|
||||
)
|
||||
|
||||
self._initialize_weights()
|
||||
|
||||
def _initialize_weights(self):
|
||||
nn.init.trunc_normal_(self.caption_embedder[1].weight, std=0.02)
|
||||
nn.init.zeros_(self.caption_embedder[1].bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
instruction_hidden_states: torch.Tensor,
|
||||
dtype: torch.dtype,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
timestep_proj = self.time_proj(timestep).to(dtype=dtype)
|
||||
time_embed = self.timestep_embedder(timestep_proj)
|
||||
caption_embed = self.caption_embedder(instruction_hidden_states)
|
||||
return time_embed, caption_embed
|
||||
@@ -0,0 +1,112 @@
|
||||
# Vendored from the Boogu-Image repository (boogu/models/embeddings.py).
|
||||
# Original work: Copyright 2024 The HuggingFace Team. Apache-2.0.
|
||||
#
|
||||
# Only the pieces the Boogu transformer actually needs are kept here:
|
||||
# ``TimestepEmbedding`` and ``apply_rotary_emb``.
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from diffusers.models.activations import get_activation
|
||||
from torch import nn
|
||||
|
||||
|
||||
class TimestepEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
time_embed_dim: int,
|
||||
act_fn: str = "silu",
|
||||
out_dim: int = None,
|
||||
post_act_fn: Optional[str] = None,
|
||||
cond_proj_dim=None,
|
||||
sample_proj_bias=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
|
||||
|
||||
if cond_proj_dim is not None:
|
||||
self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
|
||||
else:
|
||||
self.cond_proj = None
|
||||
|
||||
self.act = get_activation(act_fn)
|
||||
|
||||
if out_dim is not None:
|
||||
time_embed_dim_out = out_dim
|
||||
else:
|
||||
time_embed_dim_out = time_embed_dim
|
||||
self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
|
||||
|
||||
if post_act_fn is None:
|
||||
self.post_act = None
|
||||
else:
|
||||
self.post_act = get_activation(post_act_fn)
|
||||
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self):
|
||||
nn.init.normal_(self.linear_1.weight, std=0.02)
|
||||
nn.init.zeros_(self.linear_1.bias)
|
||||
nn.init.normal_(self.linear_2.weight, std=0.02)
|
||||
nn.init.zeros_(self.linear_2.bias)
|
||||
|
||||
def forward(self, sample, condition=None):
|
||||
if condition is not None:
|
||||
sample = sample + self.cond_proj(condition)
|
||||
sample = self.linear_1(sample)
|
||||
|
||||
if self.act is not None:
|
||||
sample = self.act(sample)
|
||||
|
||||
sample = self.linear_2(sample)
|
||||
|
||||
if self.post_act is not None:
|
||||
sample = self.post_act(sample)
|
||||
return sample
|
||||
|
||||
|
||||
def apply_rotary_emb(
|
||||
x: torch.Tensor,
|
||||
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
|
||||
use_real: bool = True,
|
||||
use_real_unbind_dim: int = -1,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply rotary embeddings to input tensors using the given frequency tensor.
|
||||
|
||||
Boogu always calls this with ``use_real=False`` (the Lumina-style complex
|
||||
path): ``freqs_cis`` is a complex tensor and ``x`` is reinterpreted as
|
||||
complex, multiplied, and returned as real.
|
||||
"""
|
||||
if use_real:
|
||||
cos, sin = freqs_cis # [S, D]
|
||||
cos = cos[None, None]
|
||||
sin = sin[None, None]
|
||||
cos, sin = cos.to(x.device), sin.to(x.device)
|
||||
|
||||
if use_real_unbind_dim == -1:
|
||||
# Used for flux, cogvideox, hunyuan-dit
|
||||
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)
|
||||
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
||||
elif use_real_unbind_dim == -2:
|
||||
# Used for Stable Audio, Boogu and CogView4
|
||||
x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)
|
||||
x_rotated = torch.cat([-x_imag, x_real], dim=-1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2."
|
||||
)
|
||||
|
||||
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
|
||||
|
||||
return out
|
||||
else:
|
||||
# used for lumina / boogu
|
||||
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)
|
||||
|
||||
return x_out.type_as(x)
|
||||
231
extensions_built_in/diffusion_models/boogu_image/src/pipeline.py
Normal file
231
extensions_built_in/diffusion_models/boogu_image/src/pipeline.py
Normal file
@@ -0,0 +1,231 @@
|
||||
"""Packing / sampling helpers for Boogu-Image (base T2I).
|
||||
|
||||
This module glues the Qwen3-VL instruction features and the image latents into
|
||||
the call the Boogu transformer expects, and provides a minimal flow-matching
|
||||
sampler used to render preview images during training.
|
||||
|
||||
Time convention
|
||||
---------------
|
||||
Boogu's native flow time is ``t in [0, 1]`` with ``t=0`` pure noise and ``t=1``
|
||||
clean; the transformer predicts ``clean - noise``. ai-toolkit's scheduler uses
|
||||
the opposite convention (``t=1`` noise, velocity ``noise - clean``). The
|
||||
conversion lives in ``BooguImageModel.get_noise_prediction``; this sampler runs
|
||||
entirely in Boogu's native domain via :func:`run_boogu_transformer`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from .transformer import BooguImageTransformer2DModel
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Instruction feature padding.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def pad_instruction_features(
|
||||
features_list: List[torch.Tensor],
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Right-pad per-sample ``(L_i, D)`` instruction features into a batch.
|
||||
|
||||
Captions are stored per-sample at their natural length and only padded to the
|
||||
batch max here, right before the model call. Returns ``(features (B, L, D),
|
||||
attention_mask (B, L))`` with the mask 1 for real tokens, 0 for padding.
|
||||
"""
|
||||
lengths = [f.shape[0] for f in features_list]
|
||||
max_len = max(lengths)
|
||||
dim = features_list[0].shape[-1]
|
||||
batch_size = len(features_list)
|
||||
|
||||
features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype)
|
||||
mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device)
|
||||
for i, f in enumerate(features_list):
|
||||
n = f.shape[0]
|
||||
features[i, :n] = f.to(device, dtype)
|
||||
mask[i, :n] = 1
|
||||
return features, mask
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Time-shift schedule (mirrors the released Boogu base scheduler: v1 shift).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _lin_shift(
|
||||
num_tokens: float,
|
||||
x1: float = 256.0,
|
||||
y1: float = 0.5,
|
||||
x2: float = 4096.0,
|
||||
y2: float = 1.15,
|
||||
) -> float:
|
||||
"""Linear token-count -> mu mapping (Boogu base_shift/max_shift defaults)."""
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return m * num_tokens + b
|
||||
|
||||
|
||||
def boogu_time_schedule(
|
||||
num_steps: int,
|
||||
num_patch_tokens: int,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Boogu native-domain timesteps (0=noise .. 1=clean) with v1 time shift.
|
||||
|
||||
Returns a length ``num_steps + 1`` tensor; the trailing ``1.0`` is the clean
|
||||
endpoint, matching the ``_timesteps`` tail in the reference scheduler.
|
||||
"""
|
||||
t_arr = np.linspace(0.0, 1.0, num_steps + 1, dtype=np.float32)[:-1]
|
||||
|
||||
mu = _lin_shift(max(1, int(num_patch_tokens)))
|
||||
eps = 1e-8
|
||||
t1 = np.clip(1.0 - t_arr, eps, 1.0 - eps)
|
||||
num = math.exp(mu)
|
||||
denom = num + (1.0 / t1 - 1.0)
|
||||
t_arr = (1.0 - num / denom).astype(np.float32)
|
||||
|
||||
times = np.concatenate([t_arr, np.ones(1, dtype=np.float32)])
|
||||
return torch.from_numpy(times).to(device=device, dtype=torch.float32)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transformer call (Boogu native time domain).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_boogu_transformer(
|
||||
transformer: BooguImageTransformer2DModel,
|
||||
latents: torch.Tensor, # (B, 16, H, W)
|
||||
boogu_t: torch.Tensor, # (B,) in [0, 1], 0=noise, 1=clean
|
||||
instruction_features: torch.Tensor, # (B, L, instruction_feat_dim)
|
||||
instruction_mask: torch.Tensor, # (B, L) 1 for real tokens
|
||||
freqs_cis, # precomputed per-axis rotary tables
|
||||
ref_image_hidden_states=None, # edit/TI2I: List[List[(16, H, W)]] per batch item
|
||||
) -> torch.Tensor:
|
||||
"""Run the transformer and return the raw model velocity (``clean - noise``).
|
||||
|
||||
Shapes pass straight through: the prediction comes back as ``(B, 16, H, W)``
|
||||
in the same latent layout as ``latents``. ``ref_image_hidden_states`` stays
|
||||
``None`` for the base T2I model and carries reference-image VAE latents for
|
||||
the edit (TI2I) model.
|
||||
"""
|
||||
out = transformer(
|
||||
hidden_states=latents,
|
||||
timestep=boogu_t,
|
||||
instruction_hidden_states=instruction_features,
|
||||
freqs_cis=freqs_cis,
|
||||
instruction_attention_mask=instruction_mask,
|
||||
ref_image_hidden_states=ref_image_hidden_states,
|
||||
return_dict=False,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Minimal sampling pipeline (for training previews).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BooguImagePipeline:
|
||||
"""Lightweight flow-matching sampler used by ai-toolkit's preview generation."""
|
||||
|
||||
def __init__(self, model):
|
||||
# ``model`` is the BooguImageModel so we can reuse its encode/decode and
|
||||
# latent helpers without duplicating state.
|
||||
self.model = model
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.model.device_torch
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
conditional_embeds,
|
||||
unconditional_embeds,
|
||||
height: int = 1024,
|
||||
width: int = 1024,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 4.0,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
ref_latents=None, # edit/TI2I: List[List[(16, H, W)]] reference VAE latents
|
||||
**kwargs,
|
||||
) -> List[Image.Image]:
|
||||
model = self.model
|
||||
device = model.device_torch
|
||||
dtype = model.torch_dtype
|
||||
transformer = model.transformer
|
||||
patch = model.patch_size
|
||||
ae_scale = model.vae_scale_factor # 8
|
||||
|
||||
latent_channels = transformer.config.in_channels
|
||||
h_lat = height // ae_scale
|
||||
w_lat = width // ae_scale
|
||||
num_patch_tokens = (h_lat // patch) * (w_lat // patch)
|
||||
|
||||
freqs_cis = model.get_freqs_cis()
|
||||
|
||||
do_cfg = guidance_scale > 1.0
|
||||
|
||||
if latents is None:
|
||||
shape = (1, latent_channels, h_lat, w_lat)
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=torch.float32
|
||||
)
|
||||
# In Boogu's domain t=0 is pure noise, so the initial latent IS the noise.
|
||||
latents = latents.to(device, dtype=torch.float32)
|
||||
|
||||
cond_feats, cond_mask = pad_instruction_features(
|
||||
conditional_embeds.text_embeds, device, dtype
|
||||
)
|
||||
if do_cfg:
|
||||
uncond_feats, uncond_mask = pad_instruction_features(
|
||||
unconditional_embeds.text_embeds, device, dtype
|
||||
)
|
||||
|
||||
times = boogu_time_schedule(num_inference_steps, num_patch_tokens, device)
|
||||
|
||||
for t, t_next in zip(times[:-1], times[1:]):
|
||||
boogu_t = t.expand(latents.shape[0])
|
||||
v_cond = run_boogu_transformer(
|
||||
transformer,
|
||||
latents.to(dtype),
|
||||
boogu_t,
|
||||
cond_feats,
|
||||
cond_mask,
|
||||
freqs_cis,
|
||||
ref_image_hidden_states=ref_latents,
|
||||
)
|
||||
if do_cfg:
|
||||
v_uncond = run_boogu_transformer(
|
||||
transformer,
|
||||
latents.to(dtype),
|
||||
boogu_t,
|
||||
uncond_feats,
|
||||
uncond_mask,
|
||||
freqs_cis,
|
||||
ref_image_hidden_states=ref_latents,
|
||||
)
|
||||
v = v_uncond + guidance_scale * (v_cond - v_uncond)
|
||||
else:
|
||||
v = v_cond
|
||||
latents = latents + v.to(torch.float32) * (t_next - t)
|
||||
|
||||
images = model.decode_latents(latents, device=device, dtype=dtype)
|
||||
images = images.float().clamp(-1.0, 1.0)
|
||||
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
|
||||
images = images.permute(0, 2, 3, 1).cpu().numpy()
|
||||
return [Image.fromarray(arr) for arr in images]
|
||||
244
extensions_built_in/diffusion_models/boogu_image/src/rope.py
Normal file
244
extensions_built_in/diffusion_models/boogu_image/src/rope.py
Normal file
@@ -0,0 +1,244 @@
|
||||
# Vendored from the Boogu-Image repository (boogu/models/transformers/rope.py).
|
||||
# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0.
|
||||
#
|
||||
# Only the double-stream rotary embedder (the one the transformer uses) and the
|
||||
# ``get_freqs_cis`` precompute helper are kept. The MPS-specific branch is
|
||||
# preserved verbatim.
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.models.embeddings import get_1d_rotary_pos_embed
|
||||
from einops import repeat
|
||||
|
||||
|
||||
def get_freqs_cis(
|
||||
axes_dim: Tuple[int, int, int], axes_lens: Tuple[int, int, int], theta: int
|
||||
) -> List[torch.Tensor]:
|
||||
"""Precompute the per-axis rotary frequency tables (done once per resolution)."""
|
||||
freqs_cis = []
|
||||
freqs_dtype = torch.float32 if 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
|
||||
|
||||
|
||||
class BooguImageDoubleStreamRotaryPosEmbed(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
theta: int,
|
||||
axes_dim: Tuple[int, int, int],
|
||||
axes_lens: Tuple[int, int, int] = (300, 512, 512),
|
||||
patch_size: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
self.axes_lens = axes_lens
|
||||
self.patch_size = patch_size
|
||||
|
||||
@staticmethod
|
||||
def get_freqs_cis(
|
||||
axes_dim: Tuple[int, int, int], axes_lens: Tuple[int, int, int], theta: int
|
||||
) -> List[torch.Tensor]:
|
||||
return get_freqs_cis(axes_dim, axes_lens, theta)
|
||||
|
||||
def _get_freqs_cis(self, freqs_cis, 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 = 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)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
freqs_cis,
|
||||
attention_mask,
|
||||
l_effective_ref_img_len,
|
||||
l_effective_img_len,
|
||||
ref_img_sizes,
|
||||
img_sizes,
|
||||
device,
|
||||
):
|
||||
batch_size = len(attention_mask)
|
||||
p = self.patch_size
|
||||
|
||||
encoder_seq_len = attention_mask.shape[1]
|
||||
l_effective_cap_len = attention_mask.sum(dim=1).tolist()
|
||||
|
||||
seq_lengths = [
|
||||
cap_len + sum(ref_img_len) + img_len
|
||||
for cap_len, ref_img_len, img_len in zip(
|
||||
l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len
|
||||
)
|
||||
]
|
||||
|
||||
max_seq_len = max(seq_lengths)
|
||||
max_ref_img_len = max(
|
||||
[sum(ref_img_len) for ref_img_len in l_effective_ref_img_len]
|
||||
)
|
||||
max_img_len = max(l_effective_img_len)
|
||||
|
||||
# Create position IDs
|
||||
position_ids = torch.zeros(
|
||||
batch_size, max_seq_len, 3, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
for i, (cap_seq_len, seq_len) in enumerate(
|
||||
zip(l_effective_cap_len, seq_lengths)
|
||||
):
|
||||
# add text position ids
|
||||
position_ids[i, :cap_seq_len] = repeat(
|
||||
torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
|
||||
)
|
||||
|
||||
pe_shift = cap_seq_len
|
||||
pe_shift_len = cap_seq_len
|
||||
|
||||
if ref_img_sizes[i] is not None:
|
||||
for ref_img_size, ref_img_len in zip(
|
||||
ref_img_sizes[i], l_effective_ref_img_len[i]
|
||||
):
|
||||
H, W = ref_img_size
|
||||
ref_H_tokens, ref_W_tokens = H // p, W // p
|
||||
assert ref_H_tokens * ref_W_tokens == ref_img_len
|
||||
|
||||
row_ids = repeat(
|
||||
torch.arange(ref_H_tokens, dtype=torch.int32, device=device),
|
||||
"h -> h w",
|
||||
w=ref_W_tokens,
|
||||
).flatten()
|
||||
col_ids = repeat(
|
||||
torch.arange(ref_W_tokens, dtype=torch.int32, device=device),
|
||||
"w -> h w",
|
||||
h=ref_H_tokens,
|
||||
).flatten()
|
||||
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = (
|
||||
pe_shift
|
||||
)
|
||||
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = (
|
||||
row_ids
|
||||
)
|
||||
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = (
|
||||
col_ids
|
||||
)
|
||||
|
||||
pe_shift += max(ref_H_tokens, ref_W_tokens)
|
||||
pe_shift_len += ref_img_len
|
||||
|
||||
H, W = img_sizes[i]
|
||||
H_tokens, W_tokens = H // p, W // p
|
||||
assert H_tokens * W_tokens == l_effective_img_len[i]
|
||||
|
||||
row_ids = repeat(
|
||||
torch.arange(H_tokens, dtype=torch.int32, device=device),
|
||||
"h -> h w",
|
||||
w=W_tokens,
|
||||
).flatten()
|
||||
col_ids = repeat(
|
||||
torch.arange(W_tokens, dtype=torch.int32, device=device),
|
||||
"w -> h w",
|
||||
h=H_tokens,
|
||||
).flatten()
|
||||
|
||||
assert pe_shift_len + l_effective_img_len[i] == seq_len
|
||||
position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
|
||||
position_ids[i, pe_shift_len:seq_len, 1] = row_ids
|
||||
position_ids[i, pe_shift_len:seq_len, 2] = col_ids
|
||||
|
||||
# Get combined rotary embeddings
|
||||
freqs_cis = self._get_freqs_cis(freqs_cis, position_ids)
|
||||
|
||||
# create separate rotary embeddings for captions and images
|
||||
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,
|
||||
)
|
||||
|
||||
# Calculate combined image sequence lengths (ref_img + img) for each sample
|
||||
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)
|
||||
|
||||
# Create combined image rotary embeddings
|
||||
combined_img_freqs_cis = torch.zeros(
|
||||
batch_size,
|
||||
max_combined_img_len,
|
||||
freqs_cis.shape[-1],
|
||||
device=device,
|
||||
dtype=freqs_cis.dtype,
|
||||
)
|
||||
|
||||
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 image rotary embeddings: ref_img + img (same order as img_patch_embed_and_refine)
|
||||
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,
|
||||
]
|
||||
)
|
||||
|
||||
return (
|
||||
cap_freqs_cis,
|
||||
ref_img_freqs_cis,
|
||||
img_freqs_cis,
|
||||
freqs_cis,
|
||||
l_effective_cap_len,
|
||||
seq_lengths,
|
||||
combined_img_freqs_cis,
|
||||
combined_img_seq_lengths,
|
||||
)
|
||||
1174
extensions_built_in/diffusion_models/boogu_image/src/transformer.py
Normal file
1174
extensions_built_in/diffusion_models/boogu_image/src/transformer.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user