mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-09-21 05:27:57 -05:00
Support SenseNova U1.5 (CORE-411) (#15922)
This commit is contained in:
T8star-Aix
@@ -0,0 +1 @@
|
||||
"""SenseNova U1.5 model implementation."""
|
||||
@@ -0,0 +1,135 @@
|
||||
import torch
|
||||
|
||||
|
||||
IMAGE_CONTEXT_ID = 151669
|
||||
IMAGE_START_ID = 151670
|
||||
IMAGE_END_ID = 151671
|
||||
IM_START_ID = 151644
|
||||
IM_END_ID = 151645
|
||||
USER_ID = 872
|
||||
ASSISTANT_ID = 77091
|
||||
NEWLINE_ID = 198
|
||||
IMAGE_LABEL_ID = 1906
|
||||
HYPHEN_ID = 12
|
||||
DIGIT_ZERO_ID = 15
|
||||
COLON_ID = 25
|
||||
|
||||
|
||||
def preprocess_reference(image):
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
image = image[:, :, :, :3].movedim(-1, 1).float()
|
||||
if image.shape[1] == 0:
|
||||
image = image.new_zeros((image.shape[0], 3, *image.shape[-2:]))
|
||||
elif image.shape[1] < 3:
|
||||
repeats = (3 + image.shape[1] - 1) // image.shape[1]
|
||||
image = image.repeat(1, repeats, 1, 1)[:, :3]
|
||||
mean = image.new_tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1)
|
||||
std = image.new_tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1)
|
||||
return (image - mean) / std
|
||||
|
||||
|
||||
def split_reference_batches(images):
|
||||
references = []
|
||||
for image in images:
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
references.extend(image[index : index + 1] for index in range(image.shape[0]))
|
||||
return references
|
||||
|
||||
|
||||
def preprocess_references(images):
|
||||
return [preprocess_reference(image) for image in split_reference_batches(images)]
|
||||
|
||||
|
||||
def _image_tokens(token_height, token_width):
|
||||
return (
|
||||
[IMAGE_START_ID]
|
||||
+ [IMAGE_CONTEXT_ID] * (token_height * token_width)
|
||||
+ [IMAGE_END_ID]
|
||||
)
|
||||
|
||||
|
||||
def _image_label_tokens(index):
|
||||
digits = (DIGIT_ZERO_ID + int(digit) for digit in str(index + 1))
|
||||
return (IMAGE_LABEL_ID, HYPHEN_ID, *digits, COLON_ID)
|
||||
|
||||
|
||||
def conditioned_input_length(input_length, reference_grids, image_only=False):
|
||||
image_token_count = sum(height * width for height, width in reference_grids)
|
||||
if image_only:
|
||||
return image_token_count + 9 + 2 * len(reference_grids)
|
||||
label_count = (
|
||||
sum(len(_image_label_tokens(index)) for index in range(len(reference_grids)))
|
||||
if len(reference_grids) > 1
|
||||
else 0
|
||||
)
|
||||
return input_length + image_token_count + 3 * len(reference_grids) + label_count
|
||||
|
||||
|
||||
def condition_input_ids(input_ids, reference_grids, image_only=False):
|
||||
image_blocks = [_image_tokens(height, width) for height, width in reference_grids]
|
||||
if image_only:
|
||||
values = (
|
||||
[IM_START_ID, USER_ID, NEWLINE_ID]
|
||||
+ [token for block in image_blocks for token in block]
|
||||
+ [
|
||||
IM_END_ID,
|
||||
NEWLINE_ID,
|
||||
IM_START_ID,
|
||||
ASSISTANT_ID,
|
||||
NEWLINE_ID,
|
||||
IMAGE_START_ID,
|
||||
]
|
||||
)
|
||||
return torch.tensor([values], dtype=torch.long, device=input_ids.device)
|
||||
|
||||
values = input_ids[0].tolist()
|
||||
starts = [index for index, value in enumerate(values) if value == IM_START_ID]
|
||||
insert_at = starts[1] + 3 if len(starts) > 1 else len(values)
|
||||
inserted = []
|
||||
for index, block in enumerate(image_blocks):
|
||||
if len(image_blocks) > 1:
|
||||
inserted.extend(_image_label_tokens(index))
|
||||
inserted.extend(block)
|
||||
inserted.append(NEWLINE_ID)
|
||||
values[insert_at:insert_at] = inserted
|
||||
return torch.tensor([values], dtype=torch.long, device=input_ids.device)
|
||||
|
||||
|
||||
def thw_indexes(input_ids, reference_grids):
|
||||
values = input_ids[0]
|
||||
image_start_shift = torch.cat(
|
||||
(
|
||||
torch.zeros(1, dtype=torch.long, device=values.device),
|
||||
(values == IMAGE_START_ID).long(),
|
||||
)
|
||||
)[:-1]
|
||||
not_image = (values != IMAGE_CONTEXT_ID).long()
|
||||
time_indexes = (image_start_shift + not_image).cumsum(0) - 1
|
||||
height_indexes = torch.zeros_like(time_indexes)
|
||||
width_indexes = torch.zeros_like(time_indexes)
|
||||
selected = values == IMAGE_CONTEXT_ID
|
||||
height_positions = []
|
||||
width_positions = []
|
||||
for token_height, token_width in reference_grids:
|
||||
positions = torch.arange(
|
||||
token_height * token_width, dtype=torch.long, device=values.device
|
||||
)
|
||||
height_positions.append(positions // token_width)
|
||||
width_positions.append(positions % token_width)
|
||||
if height_positions:
|
||||
height_indexes[selected] = torch.cat(height_positions)
|
||||
width_indexes[selected] = torch.cat(width_positions)
|
||||
return torch.stack((time_indexes, height_indexes, width_indexes)).unsqueeze(0)
|
||||
|
||||
|
||||
def block_causal_mask(time_indexes, dtype=torch.float32):
|
||||
values = time_indexes[0, 0]
|
||||
length = values.shape[0]
|
||||
same_block = values[:, None] == values[None, :]
|
||||
positions = torch.arange(length, device=values.device)
|
||||
causal = positions[None, :] <= positions[:, None]
|
||||
allowed = same_block | causal
|
||||
mask = torch.zeros((1, 1, length, length), dtype=dtype, device=values.device)
|
||||
return mask.masked_fill_(~allowed[None, None], float("-inf"))
|
||||
@@ -0,0 +1,622 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import comfy.patcher_extension
|
||||
import comfy.utils
|
||||
from comfy.ldm.common_dit import pad_to_patch_size
|
||||
from comfy.ldm.flux.math import apply_rope1
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
|
||||
|
||||
from .sampling import resolution_noise_scale
|
||||
|
||||
|
||||
HIDDEN_SIZE = 4096
|
||||
INTERMEDIATE_SIZE = 12288
|
||||
NUM_LAYERS = 42
|
||||
NUM_HEADS = 32
|
||||
NUM_KV_HEADS = 8
|
||||
HEAD_DIM = 128
|
||||
MERGED_PATCH_SIZE = 32
|
||||
VOCAB_SIZE = 151936
|
||||
|
||||
|
||||
def _pad_to_merged_patch_size(value):
|
||||
height, width = value.shape[-2:]
|
||||
height_pad = max(16 - height, 0)
|
||||
width_pad = max(16 - width, 0)
|
||||
if height_pad or width_pad:
|
||||
value = F.pad(
|
||||
value,
|
||||
(0, width_pad, 0, height_pad),
|
||||
mode="replicate" if height > 0 and width > 0 else "constant",
|
||||
)
|
||||
return pad_to_patch_size(value, (MERGED_PATCH_SIZE, MERGED_PATCH_SIZE))
|
||||
|
||||
|
||||
def _generation_batch_size(total_batch, prefix_batch):
|
||||
if prefix_batch < 1 or total_batch < 1 or total_batch % prefix_batch != 0:
|
||||
raise ValueError(
|
||||
"SenseNova generation batch must be a positive multiple of the prefix batch "
|
||||
f"(generation={total_batch}, prefix={prefix_batch})"
|
||||
)
|
||||
return total_batch // prefix_batch
|
||||
|
||||
|
||||
def _match_prefix_batch(total_batch, text_input_ids, prefix_indexes, prefix_mask):
|
||||
prefix_batch = text_input_ids.shape[0]
|
||||
if prefix_batch > 0 and total_batch % prefix_batch:
|
||||
text_input_ids = comfy.utils.resize_to_batch_size(text_input_ids, total_batch)
|
||||
if prefix_indexes is not None:
|
||||
prefix_indexes = comfy.utils.resize_to_batch_size(
|
||||
prefix_indexes, total_batch
|
||||
)
|
||||
if prefix_mask is not None:
|
||||
prefix_mask = comfy.utils.resize_to_batch_size(prefix_mask, total_batch)
|
||||
return text_input_ids, prefix_indexes, prefix_mask
|
||||
|
||||
|
||||
def _expand_prefix_batch(value, generation_batch):
|
||||
"""Repeat each guidance branch's prefix KV for its generated variants."""
|
||||
if generation_batch == 1:
|
||||
return value
|
||||
prefix_batch = value.shape[0]
|
||||
return (
|
||||
value.unsqueeze(1)
|
||||
.expand(prefix_batch, generation_batch, *value.shape[1:])
|
||||
.reshape(prefix_batch * generation_batch, *value.shape[1:])
|
||||
)
|
||||
|
||||
|
||||
def _prepare_llm_rope(positions, dim, theta, device, dtype):
|
||||
frequencies = theta ** (
|
||||
-torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim
|
||||
)
|
||||
positions = positions.to(device=device, dtype=torch.float32)
|
||||
if positions.ndim == 1:
|
||||
positions = positions.unsqueeze(0)
|
||||
angles = positions.unsqueeze(-1) * frequencies
|
||||
embedding = torch.cat((angles, angles), dim=-1).unsqueeze(1)
|
||||
return embedding.cos().to(dtype), embedding.sin().to(dtype)
|
||||
|
||||
|
||||
def _prepare_mrope(indexes, device, dtype):
|
||||
return (
|
||||
_prepare_llm_rope(indexes[0], HEAD_DIM // 2, 5000000.0, device, dtype),
|
||||
_prepare_llm_rope(indexes[1], HEAD_DIM // 4, 10000.0, device, dtype),
|
||||
_prepare_llm_rope(indexes[2], HEAD_DIM // 4, 10000.0, device, dtype),
|
||||
)
|
||||
|
||||
|
||||
def _apply_llm_rope(query, key, rope):
|
||||
cosine, sine = rope
|
||||
|
||||
def rotate_half(value):
|
||||
first, second = value.chunk(2, dim=-1)
|
||||
return torch.cat((-second, first), dim=-1)
|
||||
|
||||
# Keep this split-half RoPE on the reference PyTorch formula. The
|
||||
# comfy-kitchen CUDA kernel is selected automatically on CUDA 13 builds;
|
||||
# on Blackwell it can return finite but numerically incorrect values, which
|
||||
# corrupts the generated image without raising an execution error.
|
||||
return (
|
||||
query * cosine + rotate_half(query) * sine,
|
||||
key * cosine + rotate_half(key) * sine,
|
||||
)
|
||||
|
||||
|
||||
def _apply_interleaved_rope(value, positions, theta):
|
||||
dim = value.shape[-1]
|
||||
frequencies = theta ** (
|
||||
-torch.arange(0, dim, 2, dtype=torch.float32, device=value.device) / dim
|
||||
)
|
||||
angles = (
|
||||
positions.to(device=value.device, dtype=torch.float32).unsqueeze(-1)
|
||||
* frequencies
|
||||
)
|
||||
cosine = angles.cos()
|
||||
sine = angles.sin()
|
||||
# comfy-kitchen acceleration backends use the canonical four-dimensional
|
||||
# input and six-dimensional rotation layout. SenseNova's vision patches
|
||||
# have no head axis, so add a singleton one instead of relying on the eager
|
||||
# backend's more permissive rank handling.
|
||||
rotation = torch.stack((cosine, -sine, sine, cosine), dim=-1).reshape(
|
||||
1, 1, *angles.shape, 2, 2
|
||||
)
|
||||
return apply_rope1(value.float().unsqueeze(1), rotation).squeeze(1)
|
||||
|
||||
|
||||
class VisionEmbeddings(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.patch_embedding = operations.Conv2d(
|
||||
3, 1024, kernel_size=16, stride=16, device=device, dtype=dtype
|
||||
)
|
||||
self.dense_embedding = operations.Conv2d(
|
||||
1024, HIDDEN_SIZE, kernel_size=2, stride=2, device=device, dtype=dtype
|
||||
)
|
||||
self.gelu = nn.GELU()
|
||||
|
||||
def forward(self, image):
|
||||
patches = self.gelu(self.patch_embedding(image))
|
||||
batch, channels, height, width = patches.shape
|
||||
patches = patches.flatten(2).transpose(1, 2)
|
||||
indexes = torch.arange(height * width, device=patches.device)
|
||||
x_positions = indexes % width
|
||||
y_positions = indexes // width
|
||||
first = _apply_interleaved_rope(
|
||||
patches[..., : channels // 2], x_positions, 10000.0
|
||||
)
|
||||
second = _apply_interleaved_rope(
|
||||
patches[..., channels // 2 :], y_positions, 10000.0
|
||||
)
|
||||
patches = torch.cat((first, second), dim=-1).to(image.dtype)
|
||||
patches = patches.transpose(1, 2).reshape(batch, channels, height, width)
|
||||
patches = self.dense_embedding(patches)
|
||||
return patches.flatten(2).transpose(1, 2)
|
||||
|
||||
|
||||
class VisionModel(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.embeddings = VisionEmbeddings(
|
||||
device=device, dtype=dtype, operations=operations
|
||||
)
|
||||
|
||||
def forward(self, image):
|
||||
return self.embeddings(image)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.gate_proj = operations.Linear(
|
||||
HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.up_proj = operations.Linear(
|
||||
HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.down_proj = operations.Linear(
|
||||
INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
return self.down_proj(
|
||||
F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.q_proj = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.q_proj_mot_gen = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.k_proj = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.k_proj_mot_gen = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.v_proj = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.v_proj_mot_gen = operations.Linear(
|
||||
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.o_proj = operations.Linear(
|
||||
NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
self.o_proj_mot_gen = operations.Linear(
|
||||
NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
self.q_norm = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.q_norm_mot_gen = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.q_norm_hw = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.q_norm_hw_mot_gen = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.k_norm = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.k_norm_mot_gen = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.k_norm_hw = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.k_norm_hw_mot_gen = operations.RMSNorm(
|
||||
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
def _project(self, hidden_states, rope, generation):
|
||||
batch, length, _ = hidden_states.shape
|
||||
if generation:
|
||||
query = self.q_proj_mot_gen(hidden_states).view(
|
||||
batch, length, NUM_HEADS, HEAD_DIM
|
||||
)
|
||||
key = self.k_proj_mot_gen(hidden_states).view(
|
||||
batch, length, NUM_KV_HEADS, HEAD_DIM
|
||||
)
|
||||
value = (
|
||||
self.v_proj_mot_gen(hidden_states)
|
||||
.view(batch, length, NUM_KV_HEADS, HEAD_DIM)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
query_t, query_hw = query.chunk(2, dim=-1)
|
||||
key_t, key_hw = key.chunk(2, dim=-1)
|
||||
query_t = self.q_norm_mot_gen(query_t).transpose(1, 2)
|
||||
query_hw = self.q_norm_hw_mot_gen(query_hw).transpose(1, 2)
|
||||
key_t = self.k_norm_mot_gen(key_t).transpose(1, 2)
|
||||
key_hw = self.k_norm_hw_mot_gen(key_hw).transpose(1, 2)
|
||||
else:
|
||||
query = self.q_proj(hidden_states).view(batch, length, NUM_HEADS, HEAD_DIM)
|
||||
key = self.k_proj(hidden_states).view(batch, length, NUM_KV_HEADS, HEAD_DIM)
|
||||
value = (
|
||||
self.v_proj(hidden_states)
|
||||
.view(batch, length, NUM_KV_HEADS, HEAD_DIM)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
query_t, query_hw = query.chunk(2, dim=-1)
|
||||
key_t, key_hw = key.chunk(2, dim=-1)
|
||||
query_t = self.q_norm(query_t).transpose(1, 2)
|
||||
query_hw = self.q_norm_hw(query_hw).transpose(1, 2)
|
||||
key_t = self.k_norm(key_t).transpose(1, 2)
|
||||
key_hw = self.k_norm_hw(key_hw).transpose(1, 2)
|
||||
|
||||
query_h, query_w = query_hw.chunk(2, dim=-1)
|
||||
key_h, key_w = key_hw.chunk(2, dim=-1)
|
||||
query_t, key_t = _apply_llm_rope(query_t, key_t, rope[0])
|
||||
query_h, key_h = _apply_llm_rope(query_h, key_h, rope[1])
|
||||
query_w, key_w = _apply_llm_rope(query_w, key_w, rope[2])
|
||||
query = torch.cat((query_t, query_h, query_w), dim=-1)
|
||||
key = torch.cat((key_t, key_h, key_w), dim=-1)
|
||||
return query, key, value
|
||||
|
||||
def forward_prefix(
|
||||
self, hidden_states, rope, attention_mask, transformer_options
|
||||
):
|
||||
query, key, value = self._project(hidden_states, rope, False)
|
||||
output = optimized_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
NUM_HEADS,
|
||||
mask=attention_mask,
|
||||
skip_reshape=True,
|
||||
transformer_options=transformer_options,
|
||||
enable_gqa=True,
|
||||
)
|
||||
return self.o_proj(output), key, value
|
||||
|
||||
def forward_generation(
|
||||
self, hidden_states, rope, prefix_key, prefix_value, transformer_options
|
||||
):
|
||||
query, key, value = self._project(hidden_states, rope, True)
|
||||
key = torch.cat((prefix_key, key), dim=2)
|
||||
value = torch.cat((prefix_value, value), dim=2)
|
||||
output = optimized_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
NUM_HEADS,
|
||||
mask=None,
|
||||
skip_reshape=True,
|
||||
transformer_options=transformer_options,
|
||||
enable_gqa=True,
|
||||
)
|
||||
return self.o_proj_mot_gen(output)
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(device=device, dtype=dtype, operations=operations)
|
||||
self.mlp = MLP(device=device, dtype=dtype, operations=operations)
|
||||
self.mlp_mot_gen = MLP(device=device, dtype=dtype, operations=operations)
|
||||
self.input_layernorm = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.input_layernorm_mot_gen = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.post_attention_layernorm = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.post_attention_layernorm_mot_gen = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
def forward_prefix(self, prefix, prefix_rope, prefix_mask, transformer_options):
|
||||
prefix_attention, prefix_key, prefix_value = self.self_attn.forward_prefix(
|
||||
self.input_layernorm(prefix),
|
||||
prefix_rope,
|
||||
prefix_mask,
|
||||
transformer_options,
|
||||
)
|
||||
prefix = prefix + prefix_attention
|
||||
prefix = prefix + self.mlp(self.post_attention_layernorm(prefix))
|
||||
return prefix, prefix_key, prefix_value
|
||||
|
||||
def forward_generation(
|
||||
self, image, image_rope, prefix_key, prefix_value, transformer_options
|
||||
):
|
||||
image_attention = self.self_attn.forward_generation(
|
||||
self.input_layernorm_mot_gen(image),
|
||||
image_rope,
|
||||
prefix_key,
|
||||
prefix_value,
|
||||
transformer_options,
|
||||
)
|
||||
image = image + image_attention
|
||||
image = image + self.mlp_mot_gen(self.post_attention_layernorm_mot_gen(image))
|
||||
return image
|
||||
|
||||
|
||||
class LanguageBackbone(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.embed_tokens = operations.Embedding(
|
||||
VOCAB_SIZE, HIDDEN_SIZE, padding_idx=151643, device=device, dtype=dtype
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
DecoderLayer(device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(NUM_LAYERS)
|
||||
)
|
||||
self.norm = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
self.norm_mot_gen = operations.RMSNorm(
|
||||
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
|
||||
class LanguageModel(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.model = LanguageBackbone(device=device, dtype=dtype, operations=operations)
|
||||
|
||||
|
||||
class ConvDecoder(nn.Module):
|
||||
def __init__(self, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = operations.Conv2d(
|
||||
1024, 1024, kernel_size=3, padding=1, device=device, dtype=dtype
|
||||
)
|
||||
self.act1 = nn.GELU()
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = operations.Conv2d(
|
||||
256, 192, kernel_size=3, padding=1, device=device, dtype=dtype
|
||||
)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
hidden_states = self.act1(self.conv1(self.ps1(hidden_states)))
|
||||
return self.ps3(self.conv2(self.ps2(hidden_states)))
|
||||
|
||||
|
||||
class SenseNovaU15(nn.Module):
|
||||
def __init__(
|
||||
self, image_model=None, dtype=None, device=None, operations=None, **kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
self.vision_model = VisionModel(
|
||||
device=device, dtype=dtype, operations=operations
|
||||
)
|
||||
self.language_model = LanguageModel(
|
||||
device=device, dtype=dtype, operations=operations
|
||||
)
|
||||
self.fm_modules = nn.ModuleDict(
|
||||
{
|
||||
"vision_model_mot_gen": VisionModel(
|
||||
device=device, dtype=dtype, operations=operations
|
||||
),
|
||||
"timestep_embedder": TimestepEmbedder(
|
||||
HIDDEN_SIZE, device=device, dtype=dtype, operations=operations
|
||||
),
|
||||
"fm_head": ConvDecoder(
|
||||
device=device, dtype=dtype, operations=operations
|
||||
),
|
||||
"noise_scale_embedder": TimestepEmbedder(
|
||||
HIDDEN_SIZE, device=device, dtype=dtype, operations=operations
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
def forward(self, x, timesteps, context=None, transformer_options={}, **kwargs):
|
||||
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
self._forward,
|
||||
self,
|
||||
comfy.patcher_extension.get_all_wrappers(
|
||||
comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options
|
||||
),
|
||||
).execute(x, timesteps, context, transformer_options, **kwargs)
|
||||
|
||||
def _prepare_prefix(
|
||||
self, text_input_ids, reference_images, prefix_indexes, prefix_mask
|
||||
):
|
||||
prefix = self.language_model.model.embed_tokens(text_input_ids)
|
||||
if reference_images:
|
||||
reference_embeds = [
|
||||
self.vision_model(_pad_to_merged_patch_size(reference))
|
||||
for reference in reference_images
|
||||
]
|
||||
selected = text_input_ids == 151669
|
||||
prefix = prefix.clone()
|
||||
prefix[selected] = torch.cat(reference_embeds, dim=1).reshape(
|
||||
-1, HIDDEN_SIZE
|
||||
)
|
||||
|
||||
prefix_length = text_input_ids.shape[1]
|
||||
if prefix_indexes is None:
|
||||
prefix_positions = torch.arange(
|
||||
prefix_length, dtype=torch.long, device=prefix.device
|
||||
)
|
||||
zeros = torch.zeros_like(prefix_positions)
|
||||
prefix_indexes = torch.stack((prefix_positions, zeros, zeros))
|
||||
prefix_mask = torch.full(
|
||||
(prefix_length, prefix_length),
|
||||
float("-inf"),
|
||||
dtype=prefix.dtype,
|
||||
device=prefix.device,
|
||||
).triu(1)
|
||||
prefix_time = torch.full(
|
||||
(prefix.shape[0],),
|
||||
prefix_length,
|
||||
dtype=torch.long,
|
||||
device=prefix.device,
|
||||
)
|
||||
else:
|
||||
prefix_indexes = prefix_indexes.transpose(0, 1)
|
||||
prefix_time = prefix_indexes[0].amax(dim=-1) + 1
|
||||
|
||||
return prefix, prefix_indexes, prefix_mask, prefix_time
|
||||
|
||||
def preprocess_prefix(
|
||||
self,
|
||||
text_input_ids,
|
||||
reference_images=None,
|
||||
prefix_indexes=None,
|
||||
prefix_mask=None,
|
||||
):
|
||||
prefix, prefix_indexes, prefix_mask, prefix_time = self._prepare_prefix(
|
||||
text_input_ids, reference_images, prefix_indexes, prefix_mask
|
||||
)
|
||||
prefix_keys = []
|
||||
prefix_values = []
|
||||
prefix_rope = _prepare_mrope(prefix_indexes, prefix.device, prefix.dtype)
|
||||
transformer_options = {}
|
||||
for layer_index, layer in enumerate(self.language_model.model.layers):
|
||||
transformer_options["block_index"] = layer_index
|
||||
prefix, prefix_key, prefix_value = layer.forward_prefix(
|
||||
prefix,
|
||||
prefix_rope,
|
||||
prefix_mask,
|
||||
transformer_options,
|
||||
)
|
||||
prefix_keys.append(prefix_key)
|
||||
prefix_values.append(prefix_value)
|
||||
return prefix_keys, prefix_values, prefix_time
|
||||
|
||||
def _forward(
|
||||
self,
|
||||
x,
|
||||
timesteps,
|
||||
context=None,
|
||||
transformer_options={},
|
||||
text_input_ids=None,
|
||||
reference_images=None,
|
||||
prefix_indexes=None,
|
||||
prefix_mask=None,
|
||||
prefix_keys=None,
|
||||
prefix_values=None,
|
||||
prefix_time=None,
|
||||
**kwargs,
|
||||
):
|
||||
if text_input_ids is None and prefix_keys is None:
|
||||
raise ValueError("SenseNova-U1.5 requires text conditioning")
|
||||
|
||||
original_height, original_width = x.shape[-2:]
|
||||
x = _pad_to_merged_patch_size(x)
|
||||
batch, _, height, width = x.shape
|
||||
if prefix_keys is None:
|
||||
text_input_ids, prefix_indexes, prefix_mask = _match_prefix_batch(
|
||||
batch, text_input_ids, prefix_indexes, prefix_mask
|
||||
)
|
||||
prefix_batch = text_input_ids.shape[0]
|
||||
if reference_images:
|
||||
reference_images = [
|
||||
comfy.utils.resize_to_batch_size(reference, prefix_batch)
|
||||
for reference in reference_images
|
||||
]
|
||||
else:
|
||||
reference_images = None
|
||||
else:
|
||||
prefix_batch = prefix_keys[0].shape[0]
|
||||
if prefix_batch > 0 and batch % prefix_batch:
|
||||
prefix_keys = [
|
||||
comfy.utils.resize_to_batch_size(value, batch)
|
||||
for value in prefix_keys
|
||||
]
|
||||
prefix_values = [
|
||||
comfy.utils.resize_to_batch_size(value, batch)
|
||||
for value in prefix_values
|
||||
]
|
||||
prefix_time = comfy.utils.resize_to_batch_size(prefix_time, batch)
|
||||
prefix_batch = batch
|
||||
generation_batch = _generation_batch_size(batch, prefix_batch)
|
||||
token_height = height // MERGED_PATCH_SIZE
|
||||
token_width = width // MERGED_PATCH_SIZE
|
||||
image_length = token_height * token_width
|
||||
|
||||
image = self.fm_modules["vision_model_mot_gen"](x)
|
||||
time_embedding = self.fm_modules["timestep_embedder"](timesteps, image.dtype)
|
||||
noise_scale = resolution_noise_scale(height, width) / 16.0
|
||||
scale_timesteps = torch.full_like(timesteps, noise_scale)
|
||||
time_embedding = time_embedding + self.fm_modules["noise_scale_embedder"](
|
||||
scale_timesteps, image.dtype
|
||||
)
|
||||
image = image + time_embedding[:, None, :]
|
||||
|
||||
if prefix_keys is None:
|
||||
prefix, prefix_indexes, prefix_mask, prefix_time = self._prepare_prefix(
|
||||
text_input_ids, reference_images, prefix_indexes, prefix_mask
|
||||
)
|
||||
prefix_rope = _prepare_mrope(prefix_indexes, prefix.device, prefix.dtype)
|
||||
image_time = prefix_time.repeat_interleave(generation_batch)
|
||||
|
||||
image_positions = torch.arange(image_length, dtype=torch.long, device=x.device)
|
||||
image_indexes = torch.stack(
|
||||
(
|
||||
image_time[:, None].expand(batch, image_length),
|
||||
(image_positions // token_width)[None].expand(batch, image_length),
|
||||
(image_positions % token_width)[None].expand(batch, image_length),
|
||||
)
|
||||
)
|
||||
image_rope = _prepare_mrope(image_indexes, image.device, image.dtype)
|
||||
|
||||
for layer_index, layer in enumerate(self.language_model.model.layers):
|
||||
transformer_options["block_index"] = layer_index
|
||||
if prefix_keys is None:
|
||||
prefix, prefix_key, prefix_value = layer.forward_prefix(
|
||||
prefix,
|
||||
prefix_rope,
|
||||
prefix_mask,
|
||||
transformer_options,
|
||||
)
|
||||
else:
|
||||
prefix_key = prefix_keys[layer_index]
|
||||
prefix_value = prefix_values[layer_index]
|
||||
generation_prefix_key = _expand_prefix_batch(prefix_key, generation_batch)
|
||||
generation_prefix_value = _expand_prefix_batch(
|
||||
prefix_value, generation_batch
|
||||
)
|
||||
image = layer.forward_generation(
|
||||
image,
|
||||
image_rope,
|
||||
generation_prefix_key,
|
||||
generation_prefix_value,
|
||||
transformer_options,
|
||||
)
|
||||
|
||||
image = self.language_model.model.norm_mot_gen(image)
|
||||
image = image.view(batch, token_height, token_width, HIDDEN_SIZE).permute(
|
||||
0, 3, 1, 2
|
||||
)
|
||||
predicted = self.fm_modules["fm_head"](image)
|
||||
denominator = (1.0 - timesteps).clamp_min(0.02).view(batch, 1, 1, 1)
|
||||
velocity = (x - predicted) / denominator
|
||||
return velocity[..., :original_height, :original_width]
|
||||
@@ -0,0 +1,69 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.model_sampling
|
||||
|
||||
|
||||
def time_snr_shift(shift, value):
|
||||
if shift == 1.0:
|
||||
return value
|
||||
return shift * value / (1.0 + (shift - 1.0) * value)
|
||||
|
||||
|
||||
def inverse_time_snr_shift(shift, value):
|
||||
if shift == 1.0:
|
||||
return value
|
||||
return value / (shift - (shift - 1.0) * value)
|
||||
|
||||
|
||||
def upstream_timesteps(steps, shift, device=None):
|
||||
base = torch.linspace(0.0, 1.0, steps + 1, device=device)
|
||||
return 1.0 - time_snr_shift(shift, 1.0 - base)
|
||||
|
||||
|
||||
def upstream_sigmas(steps, shift, device=None):
|
||||
return 1.0 - upstream_timesteps(steps, shift, device=device)
|
||||
|
||||
|
||||
def resolution_noise_scale(
|
||||
height, width, base_seq_len=64, noise_scale=1.0, maximum=16.0
|
||||
):
|
||||
token_height = math.ceil(height / 32)
|
||||
token_width = math.ceil(width / 32)
|
||||
scale = math.sqrt(token_height * token_width / base_seq_len) * noise_scale
|
||||
return min(scale, maximum)
|
||||
|
||||
|
||||
class SenseNovaModelSampling(
|
||||
comfy.model_sampling.ModelSamplingDiscreteFlow, comfy.model_sampling.CONST
|
||||
):
|
||||
def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):
|
||||
self.shift = shift
|
||||
self.multiplier = multiplier
|
||||
base_timesteps = torch.linspace(multiplier, 0.0, timesteps + 1)
|
||||
self.register_buffer("sigmas", self.sigma(base_timesteps))
|
||||
|
||||
def timestep(self, sigma):
|
||||
base_sigma = inverse_time_snr_shift(self.shift, sigma)
|
||||
return (1.0 - base_sigma) * self.multiplier
|
||||
|
||||
def sigma(self, timestep):
|
||||
base_sigma = 1.0 - timestep / self.multiplier
|
||||
return time_snr_shift(self.shift, base_sigma)
|
||||
|
||||
def percent_to_sigma(self, percent):
|
||||
if percent <= 0.0:
|
||||
return 1.0
|
||||
if percent >= 1.0:
|
||||
return 0.0
|
||||
return float(time_snr_shift(self.shift, 1.0 - percent))
|
||||
|
||||
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
|
||||
sigma = comfy.model_sampling.reshape_sigma(sigma, noise.ndim)
|
||||
scale = resolution_noise_scale(
|
||||
latent_image.shape[-2],
|
||||
latent_image.shape[-1],
|
||||
noise_scale=self.noise_scale,
|
||||
)
|
||||
return sigma * (scale * noise) + (1.0 - sigma) * latent_image
|
||||
@@ -76,6 +76,9 @@ import comfy.ldm.ernie.model
|
||||
import comfy.ldm.sam3.detector
|
||||
import comfy.ldm.hidream_o1.model
|
||||
from comfy.ldm.hidream_o1.conditioning import build_extra_conds
|
||||
import comfy.ldm.sensenova.conditioning
|
||||
import comfy.ldm.sensenova.model
|
||||
from comfy.ldm.sensenova.sampling import SenseNovaModelSampling, time_snr_shift
|
||||
import comfy.ldm.depth_anything_3.model
|
||||
|
||||
import comfy.model_management
|
||||
@@ -2342,6 +2345,134 @@ class HiDreamO1(BaseModel):
|
||||
out[k] = cls(v)
|
||||
return out
|
||||
|
||||
class SenseNovaSharedRegular(comfy.conds.CONDRegular):
|
||||
"""Keep the shared text/reference prefix at one copy per guidance branch."""
|
||||
|
||||
def process_cond(self, batch_size, **kwargs):
|
||||
return self._copy_with(self.cond)
|
||||
|
||||
class SenseNovaSharedList(comfy.conds.CONDList):
|
||||
def process_cond(self, batch_size, **kwargs):
|
||||
return self._copy_with(self.cond)
|
||||
|
||||
class SenseNovaU15(BaseModel):
|
||||
PATCH_SIZE = 32
|
||||
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.sensenova.model.SenseNovaU15)
|
||||
self.model_sampling = SenseNovaModelSampling(model_config)
|
||||
self.memory_usage_factor_conds = ("reference_images",)
|
||||
|
||||
def process_timestep(self, timestep, **kwargs):
|
||||
base_timestep = timestep / self.model_sampling.multiplier
|
||||
return 1.0 - time_snr_shift(self.model_sampling.shift, 1.0 - base_timestep)
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
text_input_ids = kwargs.get("text_input_ids")
|
||||
if text_input_ids is not None:
|
||||
device = kwargs["device"]
|
||||
reference_images = kwargs.get("reference_latents")
|
||||
if reference_images is not None:
|
||||
reference_images = comfy.ldm.sensenova.conditioning.preprocess_references(reference_images)
|
||||
image_only = kwargs.get("prompt_type") == "negative"
|
||||
indexes = None
|
||||
prefix_mask = None
|
||||
if reference_images:
|
||||
reference_grids = [
|
||||
(
|
||||
max(1, math.ceil(image.shape[-2] / self.PATCH_SIZE)),
|
||||
max(1, math.ceil(image.shape[-1] / self.PATCH_SIZE)),
|
||||
)
|
||||
for image in reference_images
|
||||
]
|
||||
text_input_ids = comfy.ldm.sensenova.conditioning.condition_input_ids(
|
||||
text_input_ids,
|
||||
reference_grids,
|
||||
image_only=image_only,
|
||||
)
|
||||
indexes = comfy.ldm.sensenova.conditioning.thw_indexes(text_input_ids, reference_grids)
|
||||
prefix_mask = comfy.ldm.sensenova.conditioning.block_causal_mask(
|
||||
indexes, dtype=self.get_dtype_inference()
|
||||
)
|
||||
|
||||
if kwargs.get("hooks") is None:
|
||||
dtype = self.get_dtype_inference()
|
||||
prefix_keys, prefix_values, prefix_time = (
|
||||
self.diffusion_model.preprocess_prefix(
|
||||
text_input_ids.to(device=device),
|
||||
[
|
||||
image.to(device=device, dtype=dtype)
|
||||
for image in reference_images
|
||||
]
|
||||
if reference_images
|
||||
else None,
|
||||
indexes.to(device=device) if indexes is not None else None,
|
||||
prefix_mask.to(device=device)
|
||||
if prefix_mask is not None
|
||||
else None,
|
||||
)
|
||||
)
|
||||
out["prefix_keys"] = SenseNovaSharedList(prefix_keys)
|
||||
out["prefix_values"] = SenseNovaSharedList(prefix_values)
|
||||
out["prefix_time"] = SenseNovaSharedRegular(prefix_time)
|
||||
else:
|
||||
if reference_images:
|
||||
out["prefix_indexes"] = SenseNovaSharedRegular(indexes)
|
||||
out["prefix_mask"] = SenseNovaSharedRegular(prefix_mask)
|
||||
out["reference_images"] = SenseNovaSharedList(reference_images)
|
||||
out["text_input_ids"] = SenseNovaSharedRegular(text_input_ids)
|
||||
return out
|
||||
|
||||
def extra_conds_shapes(self, **kwargs):
|
||||
images = kwargs.get("reference_latents")
|
||||
images = comfy.ldm.sensenova.conditioning.split_reference_batches(images) if images is not None else []
|
||||
reference_grids = [
|
||||
(
|
||||
max(1, math.ceil(image.shape[-3] / self.PATCH_SIZE)),
|
||||
max(1, math.ceil(image.shape[-2] / self.PATCH_SIZE)),
|
||||
)
|
||||
for image in images
|
||||
]
|
||||
reference_pixels = sum(
|
||||
height * width * self.PATCH_SIZE**2
|
||||
for height, width in reference_grids
|
||||
)
|
||||
out = {}
|
||||
if reference_pixels:
|
||||
out["reference_images"] = [1, 3, reference_pixels]
|
||||
text_input_ids = kwargs.get("text_input_ids")
|
||||
if text_input_ids is not None:
|
||||
if reference_grids:
|
||||
length = comfy.ldm.sensenova.conditioning.conditioned_input_length(
|
||||
text_input_ids.shape[1],
|
||||
reference_grids,
|
||||
image_only=kwargs.get("prompt_type") == "negative",
|
||||
)
|
||||
else:
|
||||
length = text_input_ids.shape[1]
|
||||
out["prefix_mask"] = [1, 1, length, length]
|
||||
if kwargs.get("hooks") is None:
|
||||
prefix_shape = [
|
||||
1,
|
||||
comfy.ldm.sensenova.model.NUM_KV_HEADS,
|
||||
comfy.ldm.sensenova.model.NUM_LAYERS
|
||||
* length
|
||||
* comfy.ldm.sensenova.model.HEAD_DIM,
|
||||
]
|
||||
out["prefix_keys"] = prefix_shape
|
||||
out["prefix_values"] = prefix_shape
|
||||
return out
|
||||
|
||||
def memory_required(self, input_shape, cond_shapes={}):
|
||||
memory = super().memory_required(input_shape, cond_shapes)
|
||||
dtype_size = comfy.model_management.dtype_size(self.get_dtype_inference())
|
||||
return memory + sum(
|
||||
math.prod(shape) * dtype_size
|
||||
for key in ("prefix_mask", "prefix_keys", "prefix_values")
|
||||
for shape in cond_shapes.get(key, ())
|
||||
)
|
||||
|
||||
class Chroma(Flux):
|
||||
def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.chroma.model.Chroma):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=unet_model)
|
||||
|
||||
@@ -815,6 +815,16 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
if '{}t_embedder1.mlp.0.weight'.format(key_prefix) in state_dict_keys and '{}x_embedder.proj1.weight'.format(key_prefix) in state_dict_keys: # HiDream-O1
|
||||
return {"image_model": "hidream_o1"}
|
||||
|
||||
vision_key = f"{key_prefix}fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight"
|
||||
query_key = f"{key_prefix}language_model.model.layers.0.self_attn.q_proj_mot_gen.weight"
|
||||
if (
|
||||
vision_key in state_dict
|
||||
and query_key in state_dict
|
||||
and state_dict[vision_key].shape[0] == 1024
|
||||
and state_dict[query_key].shape[0] == 4096
|
||||
): # SenseNova U1.5
|
||||
return {"image_model": "sensenova_u15"}
|
||||
|
||||
if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys: # HiDream
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "hidream"
|
||||
@@ -1301,6 +1311,13 @@ def unet_prefix_from_state_dict(state_dict):
|
||||
if any(k.startswith("detector.") for k in state_dict) and any(k.startswith("tracker.") for k in state_dict):
|
||||
return ""
|
||||
|
||||
# SenseNova checkpoints store the diffusion and language backbones at top level.
|
||||
if (
|
||||
"fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight" in state_dict
|
||||
and "language_model.model.layers.0.self_attn.q_proj_mot_gen.weight" in state_dict
|
||||
):
|
||||
return ""
|
||||
|
||||
candidates = ["model.diffusion_model.", #ldm/sgm models
|
||||
"model.model.", #audio models
|
||||
"net.", #cosmos
|
||||
|
||||
@@ -37,6 +37,7 @@ import comfy.text_encoders.longcat_image
|
||||
import comfy.text_encoders.ernie
|
||||
import comfy.text_encoders.cogvideo
|
||||
import comfy.text_encoders.hidream_o1
|
||||
import comfy.text_encoders.sensenova
|
||||
import comfy.text_encoders.pixeldit
|
||||
|
||||
from . import supported_models_base
|
||||
@@ -1720,6 +1721,44 @@ class HiDreamO1(supported_models_base.BASE):
|
||||
comfy.text_encoders.hidream_o1.HiDreamO1TE,
|
||||
)
|
||||
|
||||
class SenseNovaU15(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "sensenova_u15",
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"shift": 3.0,
|
||||
"noise_scale": 1.0,
|
||||
}
|
||||
|
||||
latent_format = latent_formats.HiDreamO1Pixel
|
||||
memory_usage_factor = 0.033
|
||||
supported_inference_dtypes = [torch.bfloat16, torch.float32]
|
||||
|
||||
vae_key_prefix = ["vae."]
|
||||
text_encoder_key_prefix = ["text_encoders."]
|
||||
|
||||
optimizations = {"fp8": False}
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
return model_base.SenseNovaU15(self, device=device)
|
||||
|
||||
def process_unet_state_dict(self, state_dict):
|
||||
state_dict.pop("language_model.lm_head.weight", None)
|
||||
return state_dict
|
||||
|
||||
def process_vae_state_dict(self, state_dict):
|
||||
return {"pixel_space_vae": torch.tensor(1.0)}
|
||||
|
||||
def process_clip_state_dict(self, state_dict):
|
||||
return {"_sensenova_te_sentinel": torch.zeros(1)}
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
return supported_models_base.ClipTarget(
|
||||
comfy.text_encoders.sensenova.SenseNovaTokenizer,
|
||||
comfy.text_encoders.sensenova.SenseNovaTextEncoder,
|
||||
)
|
||||
|
||||
class Chroma(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "chroma",
|
||||
@@ -2537,6 +2576,7 @@ models = [
|
||||
TripoSplat,
|
||||
HiDream,
|
||||
HiDreamO1,
|
||||
SenseNovaU15,
|
||||
Chroma,
|
||||
SeedVR2,
|
||||
ChromaRadiance,
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Tokenizer-only conditioning for SenseNova U1.5.
|
||||
|
||||
The language model is part of the diffusion checkpoint, so CLIP only needs to
|
||||
produce token ids. SenseNova extends the Qwen vocabulary with image-control
|
||||
tokens; their order is significant because the checkpoint embeds them by id.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from transformers import Qwen2Tokenizer
|
||||
|
||||
from comfy import sd1_clip
|
||||
|
||||
|
||||
SYSTEM_MESSAGE = (
|
||||
"You are an image generation and editing assistant that accurately understands and executes "
|
||||
"user intent.\n\nYou support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you "
|
||||
"MUST start with a <think></think> block. Put all reasoning inside the block using plain text. "
|
||||
"DO NOT include any image tags. Keep it reasonable and directly useful for producing the final "
|
||||
"image.\n\n2. Non-Think Mode:\nIf no reasoning is needed, directly produce the final image.\n\n"
|
||||
"Task Types:\n\nA. Text-to-Image Generation:\n"
|
||||
"- Generate a high-quality image based on the user's description.\n"
|
||||
"- Ensure visual clarity, semantic consistency, and completeness.\n"
|
||||
"- DO NOT introduce elements that contradict or override the user's intent.\n\n"
|
||||
"B. Image Editing:\n"
|
||||
"- Use the provided image(s) as input or reference for modification or transformation.\n"
|
||||
"- The result can be an edited image or a new image based on the reference(s).\n"
|
||||
"- Preserve all unspecified attributes unless explicitly changed.\n\n"
|
||||
"General Rules:\n"
|
||||
"- For any visible text in the image, follow the language specified for the rendered text in "
|
||||
"the user's description, not the language of the prompt. If no language is specified, use the "
|
||||
"user's input language."
|
||||
)
|
||||
|
||||
|
||||
def build_generation_prompt(text):
|
||||
return (
|
||||
f"<|im_start|>system\n{SYSTEM_MESSAGE}<|im_end|>\n"
|
||||
f"<|im_start|>user\n{text}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n<think>\n\n</think>\n\n<img>"
|
||||
)
|
||||
|
||||
|
||||
def build_unconditional_prompt():
|
||||
return "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n<img>"
|
||||
|
||||
|
||||
class SenseNovaQwen2Tokenizer:
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args, **kwargs):
|
||||
tokenizer = Qwen2Tokenizer.from_pretrained(*args, **kwargs)
|
||||
existing_special_tokens = [
|
||||
token
|
||||
for _, token in sorted(tokenizer.added_tokens_decoder.items())
|
||||
if token.special
|
||||
]
|
||||
extra_tokens = [
|
||||
"<IMG_CONTEXT>",
|
||||
"<img>",
|
||||
"</img>",
|
||||
"<quad>",
|
||||
"</quad>",
|
||||
"<ref>",
|
||||
"</ref>",
|
||||
"<box>",
|
||||
"</box>",
|
||||
"<|action_start|>",
|
||||
"<|action_end|>",
|
||||
"<|plugin|>",
|
||||
"<|interpreter|>",
|
||||
]
|
||||
extra_tokens.extend(f"<FAKE_PAD_{index}>" for index in range(254))
|
||||
tokenizer.add_special_tokens(
|
||||
{"additional_special_tokens": existing_special_tokens + extra_tokens}
|
||||
)
|
||||
return tokenizer
|
||||
|
||||
|
||||
class SenseNovaQwenTokenizer(sd1_clip.SDTokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
tokenizer_path = os.path.join(
|
||||
os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer"
|
||||
)
|
||||
super().__init__(
|
||||
tokenizer_path,
|
||||
pad_with_end=False,
|
||||
embedding_size=4096,
|
||||
embedding_key="sensenova_u15",
|
||||
tokenizer_class=SenseNovaQwen2Tokenizer,
|
||||
has_start_token=False,
|
||||
has_end_token=False,
|
||||
pad_to_max_length=False,
|
||||
max_length=99999999,
|
||||
min_length=1,
|
||||
pad_token=151643,
|
||||
tokenizer_data=tokenizer_data,
|
||||
)
|
||||
|
||||
|
||||
class SenseNovaTokenizer(sd1_clip.SD1Tokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
super().__init__(
|
||||
embedding_directory=embedding_directory,
|
||||
tokenizer_data=tokenizer_data,
|
||||
name="sensenova_u15",
|
||||
tokenizer=SenseNovaQwenTokenizer,
|
||||
)
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
|
||||
prompt = build_generation_prompt(text) if text else build_unconditional_prompt()
|
||||
tokens = super().tokenize_with_weights(
|
||||
prompt,
|
||||
return_word_ids=return_word_ids,
|
||||
disable_weights=True,
|
||||
**kwargs,
|
||||
)
|
||||
values = tokens["sensenova_u15"][0]
|
||||
values = [value for value in values if int(value[0]) != 151643]
|
||||
return {"sensenova_u15": [values]}
|
||||
|
||||
|
||||
class SenseNovaTextEncoder(torch.nn.Module):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
super().__init__()
|
||||
self.dtypes = {torch.float32}
|
||||
self.disable_offload = True
|
||||
self.device = torch.device("cpu") if device is None else torch.device(device)
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
pairs = token_weight_pairs["sensenova_u15"][0]
|
||||
input_ids = torch.tensor([[int(value[0]) for value in pairs]], dtype=torch.long)
|
||||
return (
|
||||
input_ids.unsqueeze(-1).to(torch.float32),
|
||||
None,
|
||||
{"text_input_ids": input_ids},
|
||||
)
|
||||
|
||||
def load_sd(self, sd):
|
||||
return []
|
||||
|
||||
def get_sd(self):
|
||||
return {}
|
||||
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
def set_clip_options(self, options):
|
||||
pass
|
||||
@@ -8,6 +8,9 @@ import node_helpers
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
REFERENCE_IMAGE_INPUT_SLOTS = 100
|
||||
|
||||
|
||||
class EmptyHiDreamO1LatentImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
@@ -40,8 +43,6 @@ class EmptyHiDreamO1LatentImage(io.ComfyNode):
|
||||
|
||||
|
||||
class HiDreamO1ReferenceImages(io.ComfyNode):
|
||||
"""Attach reference images to both positive and negative conditioning."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -49,9 +50,9 @@ class HiDreamO1ReferenceImages(io.ComfyNode):
|
||||
display_name="HiDream-O1 Reference Images",
|
||||
category="model/conditioning/hidream",
|
||||
description=(
|
||||
"Attach 1-10 reference images to conditioning, one for edit instruction"
|
||||
"or multiple for subject-driven personalization."
|
||||
"Attach ordered reference images to positive and negative conditioning."
|
||||
),
|
||||
search_aliases=["sensenova reference images"],
|
||||
inputs=[
|
||||
io.Conditioning.Input(id="positive"),
|
||||
io.Conditioning.Input(id="negative"),
|
||||
@@ -59,11 +60,14 @@ class HiDreamO1ReferenceImages(io.ComfyNode):
|
||||
"images",
|
||||
template=io.Autogrow.TemplateNames(
|
||||
io.Image.Input("image"),
|
||||
names=[f"image_{i}" for i in range(1, 11)],
|
||||
min=1,
|
||||
),
|
||||
tooltip=("Reference images. 1 image = instruction edit; 2-10 images = multi reference."
|
||||
names=[
|
||||
f"image_{index}"
|
||||
for index in range(1, REFERENCE_IMAGE_INPUT_SLOTS + 1)
|
||||
],
|
||||
min=0,
|
||||
),
|
||||
optional=True,
|
||||
tooltip="Reference images are used in numeric socket order.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
@@ -73,10 +77,31 @@ class HiDreamO1ReferenceImages(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, *, positive, negative, images: io.Autogrow.Type) -> io.NodeOutput:
|
||||
refs = [images[f"image_{i}"] for i in range(1, 11) if f"image_{i}" in images]
|
||||
positive = node_helpers.conditioning_set_values(positive, {"reference_latents": refs}, append=True)
|
||||
negative = node_helpers.conditioning_set_values(negative, {"reference_latents": refs}, append=True)
|
||||
def execute(
|
||||
cls, *, positive, negative, images: io.Autogrow.Type = None
|
||||
) -> io.NodeOutput:
|
||||
images = images or {}
|
||||
ordered_names = [
|
||||
f"image_{index}"
|
||||
for index in range(1, REFERENCE_IMAGE_INPUT_SLOTS + 1)
|
||||
if f"image_{index}" in images
|
||||
]
|
||||
known_names = set(ordered_names)
|
||||
refs = [images[name] for name in ordered_names]
|
||||
refs.extend(
|
||||
image for name, image in images.items() if name not in known_names
|
||||
)
|
||||
if not refs:
|
||||
return io.NodeOutput(positive, negative)
|
||||
positive = node_helpers.conditioning_set_values(
|
||||
positive, {"reference_latents": refs}, append=True
|
||||
)
|
||||
negative = node_helpers.conditioning_set_values(
|
||||
negative, {"prompt_type": "negative"}
|
||||
)
|
||||
negative = node_helpers.conditioning_set_values(
|
||||
negative, {"reference_latents": refs}, append=True
|
||||
)
|
||||
return io.NodeOutput(positive, negative)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy.ldm.sensenova.sampling import SenseNovaModelSampling
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
class SenseNovaSamplingOptions(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="SenseNovaSamplingOptions",
|
||||
display_name="SenseNova Sampling Options",
|
||||
category="model/patch/sensenova",
|
||||
description="Set the SenseNova flow shift.",
|
||||
inputs=[
|
||||
io.Model.Input(id="model"),
|
||||
io.Float.Input(id="shift", default=3.0, step=0.01),
|
||||
],
|
||||
outputs=[io.Model.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, *, model, shift: float) -> io.NodeOutput:
|
||||
patched = model.clone()
|
||||
model_sampling = SenseNovaModelSampling(patched.model.model_config)
|
||||
model_sampling.set_parameters(shift=shift)
|
||||
patched.add_object_patch("model_sampling", model_sampling)
|
||||
return io.NodeOutput(patched)
|
||||
|
||||
|
||||
class SenseNovaExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
SenseNovaSamplingOptions,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> SenseNovaExtension:
|
||||
return SenseNovaExtension()
|
||||
@@ -2521,6 +2521,7 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_void.py",
|
||||
"nodes_wandancer.py",
|
||||
"nodes_hidream_o1.py",
|
||||
"nodes_sensenova.py",
|
||||
"nodes_save_3d.py",
|
||||
"nodes_mesh_io.py",
|
||||
"nodes_moge.py",
|
||||
|
||||
@@ -0,0 +1,657 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
args.cpu = True
|
||||
|
||||
from comfy import model_base, model_detection
|
||||
import comfy.latent_formats
|
||||
import comfy.sample
|
||||
import nodes
|
||||
from comfy.ldm.sensenova import model as sensenova_model
|
||||
from comfy.ldm.sensenova.conditioning import (
|
||||
block_causal_mask,
|
||||
condition_input_ids,
|
||||
conditioned_input_length,
|
||||
preprocess_references,
|
||||
thw_indexes,
|
||||
)
|
||||
from comfy.ldm.sensenova.model import _match_prefix_batch, _pad_to_merged_patch_size
|
||||
from comfy.ldm.sensenova.sampling import (
|
||||
SenseNovaModelSampling,
|
||||
resolution_noise_scale,
|
||||
upstream_sigmas,
|
||||
)
|
||||
from comfy.text_encoders.sensenova import SenseNovaTokenizer
|
||||
from comfy_extras.nodes_hidream_o1 import HiDreamO1ReferenceImages
|
||||
from comfy_extras.nodes_sensenova import SenseNovaSamplingOptions
|
||||
|
||||
|
||||
def _minimal_state_dict():
|
||||
return {
|
||||
"fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight": torch.empty(
|
||||
1024, 3, 16, 16, device="meta"
|
||||
),
|
||||
"language_model.model.layers.0.self_attn.q_proj_mot_gen.weight": torch.empty(
|
||||
4096, 4096, device="meta"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _generation_input_ids():
|
||||
return torch.tensor(
|
||||
[
|
||||
[
|
||||
151644,
|
||||
8948,
|
||||
198,
|
||||
1,
|
||||
151645,
|
||||
198,
|
||||
151644,
|
||||
872,
|
||||
198,
|
||||
2,
|
||||
151645,
|
||||
198,
|
||||
151644,
|
||||
77091,
|
||||
198,
|
||||
151670,
|
||||
]
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
|
||||
def _tokenize_generation_prompt(text):
|
||||
tokenizer = SenseNovaTokenizer()
|
||||
values = tokenizer.tokenize_with_weights(text)["sensenova_u15"][0]
|
||||
return torch.tensor([[int(value[0]) for value in values]], dtype=torch.long)
|
||||
|
||||
|
||||
def test_sensenova_top_level_checkpoint_detection():
|
||||
state_dict = _minimal_state_dict()
|
||||
|
||||
assert model_detection.unet_prefix_from_state_dict(state_dict) == ""
|
||||
assert model_detection.detect_unet_config(state_dict, "") == {
|
||||
"image_model": "sensenova_u15"
|
||||
}
|
||||
assert (
|
||||
type(model_detection.model_config_from_unet(state_dict, "")).__name__
|
||||
== "SenseNovaU15"
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_detection_rejects_incompatible_dimensions():
|
||||
state_dict = _minimal_state_dict()
|
||||
state_dict["language_model.model.layers.0.self_attn.q_proj_mot_gen.weight"] = (
|
||||
torch.empty(2048, 2048, device="meta")
|
||||
)
|
||||
|
||||
assert model_detection.detect_unet_config(state_dict, "") is None
|
||||
|
||||
|
||||
def test_sensenova_model_config_builds_pixel_space_outputs():
|
||||
model_config = model_detection.model_config_from_unet(_minimal_state_dict(), "")
|
||||
state_dict = {
|
||||
"language_model.lm_head.weight": torch.empty(1),
|
||||
"kept": torch.empty(1),
|
||||
}
|
||||
|
||||
processed = model_config.process_unet_state_dict(state_dict)
|
||||
assert set(processed) == {"kept"}
|
||||
assert torch.equal(processed["kept"], state_dict["kept"])
|
||||
assert "pixel_space_vae" in model_config.process_vae_state_dict({})
|
||||
assert "_sensenova_te_sentinel" in model_config.process_clip_state_dict({})
|
||||
|
||||
|
||||
def test_sensenova_sampling_matches_upstream_schedule_and_resolution_scale():
|
||||
config = SimpleNamespace(sampling_settings={"shift": 3.0, "noise_scale": 1.0})
|
||||
sampling = SenseNovaModelSampling(config)
|
||||
|
||||
expected = upstream_sigmas(50, 3.0)
|
||||
actual = sampling.sigma(torch.linspace(0.0, 1000.0, 51))
|
||||
assert torch.allclose(actual, expected)
|
||||
assert sampling.percent_to_sigma(0.0) == 1.0
|
||||
assert sampling.percent_to_sigma(1.0) == 0.0
|
||||
assert resolution_noise_scale(2048, 2048) == 8.0
|
||||
assert resolution_noise_scale(4096, 4096) == 16.0
|
||||
|
||||
scaled_sampling = SenseNovaModelSampling(
|
||||
SimpleNamespace(sampling_settings={"shift": 3.0, "noise_scale": 0.5})
|
||||
)
|
||||
noise = torch.ones(1, 3, 256, 256)
|
||||
latent = torch.zeros_like(noise)
|
||||
scaled = scaled_sampling.noise_scaling(torch.ones(1), noise, latent)
|
||||
assert torch.allclose(scaled, torch.full_like(noise, 0.5))
|
||||
|
||||
|
||||
def test_shared_reference_images_append_when_chained():
|
||||
conditioning = [[torch.empty(1), {}]]
|
||||
first_image = torch.empty(1, 8, 8, 3)
|
||||
second_image = torch.empty(1, 8, 8, 3)
|
||||
|
||||
first = HiDreamO1ReferenceImages.execute(
|
||||
positive=conditioning,
|
||||
negative=conditioning,
|
||||
images={"image_1": first_image},
|
||||
)
|
||||
second = HiDreamO1ReferenceImages.execute(
|
||||
positive=first[0],
|
||||
negative=first[1],
|
||||
images={"image_1": second_image},
|
||||
)
|
||||
|
||||
references = second[0][0][1]["reference_latents"]
|
||||
assert len(references) == 2
|
||||
assert references[0] is first_image
|
||||
assert references[1] is second_image
|
||||
assert second[1][0][1]["reference_latents"] == references
|
||||
assert second[1][0][1]["prompt_type"] == "negative"
|
||||
|
||||
|
||||
def test_shared_reference_images_use_numeric_socket_order():
|
||||
conditioning = [[torch.empty(1), {}]]
|
||||
first_image = torch.empty(1, 8, 8, 3)
|
||||
second_image = torch.empty(1, 8, 8, 3)
|
||||
extra_image = torch.empty(1, 8, 8, 3)
|
||||
|
||||
output = HiDreamO1ReferenceImages.execute(
|
||||
positive=conditioning,
|
||||
negative=conditioning,
|
||||
images={
|
||||
"image_2": second_image,
|
||||
"extra_image": extra_image,
|
||||
"image_1": first_image,
|
||||
},
|
||||
)
|
||||
|
||||
references = output[0][0][1]["reference_latents"]
|
||||
assert references[0] is first_image
|
||||
assert references[1] is second_image
|
||||
assert references[2] is extra_image
|
||||
|
||||
|
||||
def test_shared_reference_images_allow_empty_inputs_and_image_batches():
|
||||
conditioning = [[torch.empty(1), {}]]
|
||||
empty = HiDreamO1ReferenceImages.execute(
|
||||
positive=conditioning,
|
||||
negative=conditioning,
|
||||
images={},
|
||||
)
|
||||
assert empty[0] is conditioning
|
||||
assert empty[1] is conditioning
|
||||
|
||||
image_batch = torch.empty(2, 8, 8, 3)
|
||||
attached = HiDreamO1ReferenceImages.execute(
|
||||
positive=conditioning,
|
||||
negative=conditioning,
|
||||
images={"image_1": image_batch},
|
||||
)
|
||||
assert attached[0][0][1]["reference_latents"] == [image_batch]
|
||||
|
||||
|
||||
def test_hidream_o1_ignores_shared_negative_marker(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def build_extra_conds(text_input_ids, noise, ref_images, target_patch_size):
|
||||
calls.append((text_input_ids, noise, ref_images, target_patch_size))
|
||||
return {
|
||||
"input_ids": text_input_ids,
|
||||
"ar_len": text_input_ids.shape[1] - 1,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(model_base, "build_extra_conds", build_extra_conds)
|
||||
model = object.__new__(model_base.HiDreamO1)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.concat_keys = ()
|
||||
input_ids = torch.tensor([[1, 2, 3]], dtype=torch.long)
|
||||
noise = torch.empty(1, 3, 64, 64)
|
||||
references = [torch.empty(1, 32, 32, 3)]
|
||||
|
||||
positive = model.extra_conds(
|
||||
text_input_ids=input_ids,
|
||||
noise=noise,
|
||||
reference_latents=references,
|
||||
)
|
||||
negative = model.extra_conds(
|
||||
text_input_ids=input_ids,
|
||||
noise=noise,
|
||||
reference_latents=references,
|
||||
prompt_type="negative",
|
||||
)
|
||||
|
||||
assert len(calls) == 2
|
||||
for call_input_ids, call_noise, call_references, call_patch_size in calls:
|
||||
assert call_input_ids is input_ids
|
||||
assert call_noise is noise
|
||||
assert call_references is references
|
||||
assert call_patch_size == 32
|
||||
assert positive.keys() == negative.keys()
|
||||
assert torch.equal(positive["input_ids"].cond, negative["input_ids"].cond)
|
||||
assert positive["ar_len"].cond == negative["ar_len"].cond
|
||||
|
||||
|
||||
def test_sensenova_reference_preprocessing_preserves_size_and_splits_batches():
|
||||
references = preprocess_references(
|
||||
[torch.rand(2, 9, 13, 1), torch.empty(1, 0, 0, 0)]
|
||||
)
|
||||
|
||||
assert len(references) == 3
|
||||
assert all(reference.shape == (1, 3, 9, 13) for reference in references[:2])
|
||||
assert references[2].shape == (1, 3, 0, 0)
|
||||
assert _pad_to_merged_patch_size(references[0]).shape == (1, 3, 32, 32)
|
||||
assert _pad_to_merged_patch_size(references[2]).shape == (1, 3, 32, 32)
|
||||
|
||||
|
||||
def test_sensenova_reference_shape_estimate_uses_padded_image_sizes():
|
||||
model = object.__new__(model_base.SenseNovaU15)
|
||||
input_ids = _generation_input_ids()
|
||||
|
||||
shapes = model.extra_conds_shapes(
|
||||
reference_latents=[torch.empty(2, 33, 65, 3)],
|
||||
text_input_ids=input_ids,
|
||||
)
|
||||
|
||||
grids = [(2, 3), (2, 3)]
|
||||
length = conditioned_input_length(input_ids.shape[1], grids)
|
||||
assert shapes["reference_images"] == [1, 3, 12288]
|
||||
assert shapes["prefix_mask"] == [1, 1, length, length]
|
||||
prefix_shape = [
|
||||
1,
|
||||
sensenova_model.NUM_KV_HEADS,
|
||||
sensenova_model.NUM_LAYERS * length * sensenova_model.HEAD_DIM,
|
||||
]
|
||||
assert shapes["prefix_keys"] == prefix_shape
|
||||
assert shapes["prefix_values"] == prefix_shape
|
||||
|
||||
negative_shapes = model.extra_conds_shapes(
|
||||
reference_latents=[torch.empty(2, 33, 65, 3)],
|
||||
text_input_ids=input_ids,
|
||||
prompt_type="negative",
|
||||
)
|
||||
negative_length = conditioned_input_length(
|
||||
input_ids.shape[1], grids, image_only=True
|
||||
)
|
||||
assert negative_shapes["prefix_mask"] == [
|
||||
1,
|
||||
1,
|
||||
negative_length,
|
||||
negative_length,
|
||||
]
|
||||
|
||||
|
||||
def test_standard_empty_latent_adapts_to_sensenova_pixel_format():
|
||||
latent = nodes.EmptyLatentImage().generate(width=64, height=96, batch_size=2)[0]
|
||||
model = SimpleNamespace(
|
||||
get_model_object=lambda name: comfy.latent_formats.HiDreamO1Pixel()
|
||||
)
|
||||
|
||||
samples = comfy.sample.fix_empty_latent_channels(
|
||||
model,
|
||||
latent["samples"],
|
||||
latent["downscale_ratio_spacial"],
|
||||
)
|
||||
|
||||
assert samples.shape == (2, 3, 96, 64)
|
||||
|
||||
|
||||
def test_sensenova_prefix_conditioning_adapts_to_mismatched_batches():
|
||||
input_ids = torch.tensor([[1], [2]])
|
||||
indexes = torch.arange(6).reshape(2, 3, 1)
|
||||
mask = torch.zeros(2, 1, 1, 1)
|
||||
|
||||
input_ids, indexes, mask = _match_prefix_batch(3, input_ids, indexes, mask)
|
||||
|
||||
assert input_ids[:, 0].tolist() == [1, 2, 2]
|
||||
assert indexes.shape == (3, 3, 1)
|
||||
assert mask.shape == (3, 1, 1, 1)
|
||||
|
||||
|
||||
def test_reference_node_and_sensenova_sampling_do_not_add_quality_limits():
|
||||
sampling_inputs = {
|
||||
input.id: input for input in SenseNovaSamplingOptions.define_schema().inputs
|
||||
}
|
||||
assert sampling_inputs["shift"].min is None
|
||||
assert sampling_inputs["shift"].max is None
|
||||
|
||||
reference_inputs = {
|
||||
input.id: input for input in HiDreamO1ReferenceImages.define_schema().inputs
|
||||
}
|
||||
images = reference_inputs["images"]
|
||||
assert images.optional
|
||||
assert images.template.min == 0
|
||||
assert len(images.template.names) == 100
|
||||
|
||||
|
||||
def test_sensenova_prefix_preprocessing_runs_each_prefix_layer_once(monkeypatch):
|
||||
calls = []
|
||||
rope_calls = []
|
||||
|
||||
prepare_mrope = sensenova_model._prepare_mrope
|
||||
|
||||
def tracked_prepare_mrope(indexes, device, dtype):
|
||||
rope_calls.append((indexes.shape, device, dtype))
|
||||
return prepare_mrope(indexes, device, dtype)
|
||||
|
||||
monkeypatch.setattr(sensenova_model, "_prepare_mrope", tracked_prepare_mrope)
|
||||
|
||||
class Layer:
|
||||
def forward_prefix(self, prefix, prefix_rope, prefix_mask, transformer_options):
|
||||
calls.append(
|
||||
(
|
||||
transformer_options["block_index"],
|
||||
tuple(axis[0].shape for axis in prefix_rope),
|
||||
prefix_mask.dtype,
|
||||
)
|
||||
)
|
||||
key = prefix[..., :1].unsqueeze(1)
|
||||
value = key + 1
|
||||
return prefix + 1, key, value
|
||||
|
||||
input_ids = torch.tensor([[1, 2, 3]])
|
||||
model = SimpleNamespace(
|
||||
language_model=SimpleNamespace(
|
||||
model=SimpleNamespace(
|
||||
embed_tokens=lambda values: torch.zeros(
|
||||
*values.shape, sensenova_model.HIDDEN_SIZE
|
||||
),
|
||||
layers=[Layer(), Layer()],
|
||||
)
|
||||
)
|
||||
)
|
||||
model._prepare_prefix = lambda *args: sensenova_model.SenseNovaU15._prepare_prefix(
|
||||
model, *args
|
||||
)
|
||||
|
||||
prefix_keys, prefix_values, prefix_time = (
|
||||
sensenova_model.SenseNovaU15.preprocess_prefix(model, input_ids)
|
||||
)
|
||||
|
||||
assert calls == [
|
||||
(
|
||||
0,
|
||||
(
|
||||
torch.Size([1, 1, 3, 64]),
|
||||
torch.Size([1, 1, 3, 32]),
|
||||
torch.Size([1, 1, 3, 32]),
|
||||
),
|
||||
torch.float32,
|
||||
),
|
||||
(
|
||||
1,
|
||||
(
|
||||
torch.Size([1, 1, 3, 64]),
|
||||
torch.Size([1, 1, 3, 32]),
|
||||
torch.Size([1, 1, 3, 32]),
|
||||
),
|
||||
torch.float32,
|
||||
),
|
||||
]
|
||||
assert rope_calls == [(torch.Size([3, 3]), torch.device("cpu"), torch.float32)]
|
||||
assert len(prefix_keys) == 2
|
||||
assert len(prefix_values) == 2
|
||||
assert prefix_keys[0].shape == (1, 1, 3, 1)
|
||||
assert prefix_time.tolist() == [3]
|
||||
|
||||
|
||||
def test_sensenova_model_base_preprocesses_prefix_conditioning():
|
||||
calls = []
|
||||
|
||||
def preprocess_prefix(input_ids, references, indexes, prefix_mask):
|
||||
calls.append((input_ids, references, indexes, prefix_mask))
|
||||
return (
|
||||
[torch.zeros(1, 1, 3, 1, dtype=torch.bfloat16)],
|
||||
[torch.ones(1, 1, 3, 1, dtype=torch.bfloat16)],
|
||||
torch.tensor([3]),
|
||||
)
|
||||
|
||||
model = object.__new__(model_base.SenseNovaU15)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.concat_keys = ()
|
||||
model.manual_cast_dtype = None
|
||||
model.diffusion_model = SimpleNamespace(
|
||||
dtype=torch.bfloat16,
|
||||
preprocess_prefix=preprocess_prefix,
|
||||
)
|
||||
input_ids = torch.tensor([[1, 2, 3]])
|
||||
|
||||
conds = model.extra_conds(
|
||||
text_input_ids=input_ids,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert calls[0][0] is input_ids
|
||||
assert calls[0][1:] == (None, None, None)
|
||||
assert "text_input_ids" not in conds
|
||||
assert conds["prefix_keys"].cond[0].dtype == torch.bfloat16
|
||||
assert conds["prefix_values"].cond[0].dtype == torch.bfloat16
|
||||
assert conds["prefix_time"].cond.tolist() == [3]
|
||||
|
||||
|
||||
def test_sensenova_uses_prompt_type_for_negative_reference_conditioning():
|
||||
calls = []
|
||||
|
||||
def preprocess_prefix(input_ids, references, indexes, prefix_mask):
|
||||
calls.append((input_ids, references, indexes, prefix_mask))
|
||||
return (
|
||||
[torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16)],
|
||||
[torch.ones(1, 1, 1, 1, dtype=torch.bfloat16)],
|
||||
torch.tensor([1]),
|
||||
)
|
||||
|
||||
model = object.__new__(model_base.SenseNovaU15)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.concat_keys = ()
|
||||
model.manual_cast_dtype = None
|
||||
model.diffusion_model = SimpleNamespace(
|
||||
dtype=torch.bfloat16,
|
||||
preprocess_prefix=preprocess_prefix,
|
||||
)
|
||||
input_ids = _generation_input_ids()
|
||||
reference = torch.rand(1, 32, 32, 3)
|
||||
|
||||
model.extra_conds(
|
||||
text_input_ids=input_ids,
|
||||
reference_latents=[reference],
|
||||
prompt_type="negative",
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
expected_ids = condition_input_ids(input_ids, [(1, 1)], image_only=True)
|
||||
assert torch.equal(calls[0][0], expected_ids)
|
||||
assert len(calls[0][1]) == 1
|
||||
assert calls[0][1][0].shape == (1, 3, 32, 32)
|
||||
assert calls[0][2].shape == (1, 3, expected_ids.shape[1])
|
||||
assert calls[0][3].shape == (1, 1, expected_ids.shape[1], expected_ids.shape[1])
|
||||
|
||||
|
||||
def test_sensenova_preprocessed_prefix_matches_raw_forward():
|
||||
class Layer:
|
||||
def forward_prefix(self, prefix, prefix_rope, prefix_mask, transformer_options):
|
||||
key = prefix[..., :1].unsqueeze(1)
|
||||
return prefix + 1, key, key + 1
|
||||
|
||||
def forward_generation(
|
||||
self, image, image_rope, prefix_key, prefix_value, transformer_options
|
||||
):
|
||||
offset = (prefix_key + prefix_value).mean(dim=(1, 2, 3))
|
||||
return image + offset[:, None, None]
|
||||
|
||||
class VisionModel:
|
||||
def __call__(self, image):
|
||||
batch, _, height, width = image.shape
|
||||
length = (height // sensenova_model.MERGED_PATCH_SIZE) * (
|
||||
width // sensenova_model.MERGED_PATCH_SIZE
|
||||
)
|
||||
return image.new_zeros(batch, length, sensenova_model.HIDDEN_SIZE)
|
||||
|
||||
class TimestepEmbedder:
|
||||
def __init__(self):
|
||||
self.shapes = []
|
||||
|
||||
def __call__(self, timesteps, dtype):
|
||||
self.shapes.append(timesteps.shape)
|
||||
return torch.zeros(
|
||||
timesteps.shape[0], sensenova_model.HIDDEN_SIZE, dtype=dtype
|
||||
)
|
||||
|
||||
class Head:
|
||||
def __call__(self, image):
|
||||
return (
|
||||
image[:, :3]
|
||||
.repeat_interleave(sensenova_model.MERGED_PATCH_SIZE, dim=2)
|
||||
.repeat_interleave(sensenova_model.MERGED_PATCH_SIZE, dim=3)
|
||||
)
|
||||
|
||||
timestep_embedder = TimestepEmbedder()
|
||||
noise_scale_embedder = TimestepEmbedder()
|
||||
backbone = SimpleNamespace(
|
||||
embed_tokens=lambda values: torch.zeros(
|
||||
*values.shape, sensenova_model.HIDDEN_SIZE
|
||||
),
|
||||
layers=[Layer(), Layer()],
|
||||
norm_mot_gen=lambda image: image,
|
||||
)
|
||||
model = SimpleNamespace(
|
||||
language_model=SimpleNamespace(model=backbone),
|
||||
fm_modules={
|
||||
"vision_model_mot_gen": VisionModel(),
|
||||
"timestep_embedder": timestep_embedder,
|
||||
"noise_scale_embedder": noise_scale_embedder,
|
||||
"fm_head": Head(),
|
||||
},
|
||||
)
|
||||
model._prepare_prefix = lambda *args: sensenova_model.SenseNovaU15._prepare_prefix(
|
||||
model, *args
|
||||
)
|
||||
input_ids = torch.tensor([[1, 2, 3]])
|
||||
image = torch.zeros(1, 3, 64, 64)
|
||||
timesteps = torch.tensor([0.5])
|
||||
|
||||
raw = sensenova_model.SenseNovaU15._forward(
|
||||
model,
|
||||
image,
|
||||
timesteps,
|
||||
text_input_ids=input_ids,
|
||||
transformer_options={},
|
||||
)
|
||||
prefix_keys, prefix_values, prefix_time = (
|
||||
sensenova_model.SenseNovaU15.preprocess_prefix(model, input_ids)
|
||||
)
|
||||
preprocessed = sensenova_model.SenseNovaU15._forward(
|
||||
model,
|
||||
image,
|
||||
timesteps,
|
||||
prefix_keys=prefix_keys,
|
||||
prefix_values=prefix_values,
|
||||
prefix_time=prefix_time,
|
||||
transformer_options={},
|
||||
)
|
||||
|
||||
assert torch.equal(raw, preprocessed)
|
||||
assert timestep_embedder.shapes == [torch.Size([1]), torch.Size([1])]
|
||||
assert noise_scale_embedder.shapes == [torch.Size([1]), torch.Size([1])]
|
||||
|
||||
|
||||
def test_sensenova_reference_tokens_and_indexes():
|
||||
input_ids = _generation_input_ids()
|
||||
grids = [(2, 3), (1, 2)]
|
||||
|
||||
conditioned = condition_input_ids(input_ids, grids)
|
||||
indexes = thw_indexes(conditioned, grids)
|
||||
|
||||
assert conditioned.shape[1] == conditioned_input_length(input_ids.shape[1], grids)
|
||||
assert torch.count_nonzero(conditioned == 151669) == 8
|
||||
assert indexes.shape == (1, 3, conditioned.shape[1])
|
||||
|
||||
|
||||
def test_sensenova_prefix_mask_matches_attention_dtype(monkeypatch):
|
||||
query = torch.empty(1, 32, 3, 128, dtype=torch.bfloat16)
|
||||
key = torch.empty(1, 8, 3, 128, dtype=torch.bfloat16)
|
||||
value = torch.empty_like(key)
|
||||
captured = {}
|
||||
|
||||
def optimized_attention(query, key, value, heads, **kwargs):
|
||||
captured.update(query=query, key=key, value=value, heads=heads, kwargs=kwargs)
|
||||
return torch.empty(1, 3, 4096, dtype=torch.bfloat16)
|
||||
|
||||
monkeypatch.setattr(sensenova_model, "optimized_attention", optimized_attention)
|
||||
attention = SimpleNamespace(
|
||||
_project=lambda hidden_states, rope, generation: (query, key, value),
|
||||
o_proj=lambda output: output,
|
||||
)
|
||||
|
||||
mask = torch.zeros(1, 1, 3, 3, dtype=torch.bfloat16)
|
||||
output, _, _ = sensenova_model.Attention.forward_prefix(
|
||||
attention,
|
||||
torch.empty(1, 3, 4096, dtype=torch.bfloat16),
|
||||
torch.empty(3, 1, 3),
|
||||
mask,
|
||||
{},
|
||||
)
|
||||
|
||||
assert output.shape == (1, 3, 4096)
|
||||
assert captured["kwargs"]["mask"] is mask
|
||||
|
||||
|
||||
def test_sensenova_prefix_mask_is_created_in_the_model_dtype():
|
||||
indexes = torch.tensor([[[0, 1, 1], [0, 0, 0], [0, 0, 0]]])
|
||||
|
||||
mask = block_causal_mask(indexes, dtype=torch.bfloat16)
|
||||
|
||||
assert mask.dtype == torch.bfloat16
|
||||
assert torch.equal(
|
||||
mask[0, 0],
|
||||
torch.tensor(
|
||||
[
|
||||
[0.0, float("-inf"), float("-inf")],
|
||||
[0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0],
|
||||
],
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_reference_tokens_allow_more_than_ten_images():
|
||||
input_ids = _tokenize_generation_prompt("test")
|
||||
grids = [(1, 1)] * 12
|
||||
|
||||
conditioned = condition_input_ids(input_ids, grids)
|
||||
indexes = thw_indexes(conditioned, grids)
|
||||
expected_text = (
|
||||
"".join(
|
||||
f"Image-{index}:<img><IMG_CONTEXT></img>\n"
|
||||
for index in range(1, 13)
|
||||
)
|
||||
+ "test"
|
||||
)
|
||||
|
||||
assert torch.equal(conditioned, _tokenize_generation_prompt(expected_text))
|
||||
assert conditioned.shape[1] == conditioned_input_length(input_ids.shape[1], grids)
|
||||
assert torch.count_nonzero(conditioned == 151669) == 12
|
||||
assert indexes.shape == (1, 3, conditioned.shape[1])
|
||||
|
||||
|
||||
def test_sensenova_reference_tokens_tolerate_nonstandard_prompt_templates():
|
||||
conditioned = condition_input_ids(torch.tensor([[1, 2]]), [(1, 1)])
|
||||
|
||||
assert torch.count_nonzero(conditioned == 151669) == 1
|
||||
|
||||
|
||||
def test_sensenova_tokenizer_control_token_ids():
|
||||
tokenizer = SenseNovaTokenizer()
|
||||
backend = tokenizer.sensenova_u15.tokenizer
|
||||
|
||||
assert len(backend) == 151936
|
||||
assert backend.convert_tokens_to_ids(
|
||||
["<IMG_CONTEXT>", "<img>", "</img>", "<FAKE_PAD_253>"]
|
||||
) == [151669, 151670, 151671, 151935]
|
||||
assert "<|im_start|>" in backend.all_special_tokens
|
||||
assert "<|vision_pad|>" in backend.all_special_tokens
|
||||
assert tokenizer.tokenize_with_weights("")["sensenova_u15"][0][-1][0] == 151670
|
||||
Reference in New Issue
Block a user