Add subject gen.
This commit is contained in:
ZHOU He
+5
-1
@@ -1,16 +1,20 @@
|
||||
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask
|
||||
from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
"BatchedMaskedReferenceGen": BatchedMaskedReferenceGen,
|
||||
"MasksToPatches": MasksToPatches,
|
||||
"ReferenceConditionCombine": ReferenceConditionCombine,
|
||||
"ConditioningMatchMask": ConditioningMatchMask,
|
||||
"GenSubject": GenSubject,
|
||||
"PreviewSubjectMask" : PreviewSubjectMask
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ReferenceConditionCombine": "Reference Condition Combine",
|
||||
"MasksToPatches": "Masks to Patches",
|
||||
"ConditioningMatchMask": "Conditioning Match Mask",
|
||||
"BatchedMaskedReferenceGen": "Batched Masked Reference Gen",
|
||||
"GenSubject": "Gen Subject",
|
||||
"PreviewSubjectMask": "Preview Subject Mask",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -400,10 +400,9 @@ class MasksToPatches:
|
||||
print(f'Patch dimensions: {h}x{w}, Total patches: {h*w}')
|
||||
return (patch_masks, h, w)
|
||||
|
||||
|
||||
class GenID:
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
class PreviewSubjectMask:
|
||||
RETURN_TYPES = ("MASK","LATENT",)
|
||||
RETURN_NAMES = ("mask","latent",)
|
||||
FUNCTION = "go"
|
||||
CATEGORY = "chara_consist"
|
||||
|
||||
@@ -411,129 +410,116 @@ class GenID:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent": ("LATENTS",),
|
||||
"target_patch_masks": ("MASK",),
|
||||
"conditions_mask": ("MASK",),
|
||||
"conditions_ref_mask": ("MASK",),
|
||||
"latent": ("LATENT",),
|
||||
"subject": ("SUBJECT",),
|
||||
},
|
||||
}
|
||||
@classmethod
|
||||
def go(cls, *, model: object, source_patch_masks: torch.Tensor, target_patch_masks: torch.Tensor,
|
||||
conditions_mask: torch.Tensor, conditions_ref_mask=torch.Tensor):
|
||||
|
||||
print(f'BatchedMaskedReferenceGen - Processing with model: {model}')
|
||||
if conditions_mask.dim() == 3:
|
||||
conditions_mask = conditions_mask[0]
|
||||
if conditions_ref_mask.dim() == 3:
|
||||
conditions_ref_mask = conditions_ref_mask[0]
|
||||
# Detect model type
|
||||
model_type = "flux"
|
||||
if model.model.__class__.__name__ == "Chroma":
|
||||
model_type = "chroma"
|
||||
print(f'Detected model type: {model_type}')
|
||||
def go(cls, latent, subject):
|
||||
return (subject[0]["front_mask"], latent, )
|
||||
|
||||
class GenSubject:
|
||||
RETURN_TYPES = ("LATENT","SUBJECT",)
|
||||
RETURN_NAMES = ("latent","subject",)
|
||||
FUNCTION = "go"
|
||||
CATEGORY = "chara_consist"
|
||||
|
||||
num_references = conditions_ref_mask.shape[0]
|
||||
num_targets = conditions_mask.shape[0] - num_references
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
|
||||
w = source_patch_masks.shape[-1]
|
||||
h = source_patch_masks.shape[1]
|
||||
patch_count = h * w
|
||||
source_patch_masks = source_patch_masks.flatten(start_dim=1)
|
||||
target_patch_masks = target_patch_masks.flatten(start_dim=1)
|
||||
reference_seq_mask = torch.log(torch.cat([conditions_mask[:num_references], torch.ones((num_references, patch_count), dtype=attn_dtype)], dim=1))
|
||||
source_seq_masks = torch.log(torch.cat([conditions_mask[:num_references+1], source_patch_masks], dim=1))
|
||||
target_seq_masks = torch.log(torch.cat([conditions_mask[num_references].unsqueeze(0).repeat(num_references+1,1), target_patch_masks], dim=1))
|
||||
sampling_step = -1
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"kv_start_step": ("INT",{"default":1,"min":0}),
|
||||
"kv_end_step": ("INT",{"default":5,"min":0}),
|
||||
"save_mask_step": ("INT",{"default":5,"min":0}),
|
||||
},
|
||||
"optional": {
|
||||
"fg_condition_mask": ("MASK",),
|
||||
}
|
||||
}
|
||||
@classmethod
|
||||
def go(cls, *, latent: torch.Tensor, kv_start_step, kv_end_step, save_mask_step, fg_condition_mask = None):
|
||||
B,L,H,W = latent["samples"].shape
|
||||
subject = [{"k": [], "v":[]}]
|
||||
def prepare_callback(model, steps, x0_output_dict=None):
|
||||
origin_callback = orig_prepare_callback(model, steps, x0_output_dict)
|
||||
def callback(step, x0, x, total_steps):
|
||||
sampling_step = step
|
||||
origin_callback(step, x0, x, total_steps)
|
||||
return callback
|
||||
|
||||
def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor:
|
||||
"""
|
||||
Implementation using PyTorch's scaled_dot_product_attention with custom masking.
|
||||
|
||||
Args:
|
||||
q: Query tensor of shape [B, H, L, D]
|
||||
k: Key tensor of shape [B, H, L, D]
|
||||
v: Value tensor of shape [B, H, L, D]
|
||||
pe: Positional encoding tensor
|
||||
attn_mask: Optional attention mask of shape [B, H, L, L]
|
||||
sampled_step = -1
|
||||
sampling_step = -1
|
||||
layer = 0
|
||||
model_type = "flux"
|
||||
w = W // model.model.diffusion_model.patch_size
|
||||
h = H // model.model.diffusion_model.patch_size
|
||||
def get_fg_attn_weights(text_len, image_start, q, k):
|
||||
attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16)
|
||||
attn_weights = attn_weights.softmax(-1).sum(1)
|
||||
if fg_condition_mask is not None:
|
||||
fg_condition_mask = rearrange(fg_condition_mask, "s -> 1 1 s")
|
||||
attn_weights = (attn_weights * fg_condition_mask).mean(1)
|
||||
mask = rearrange(attn_weights, "b (h w) -> b h w", h=h, w=w)
|
||||
return mask
|
||||
def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
|
||||
"""
|
||||
Implementation using PyTorch's scaled_dot_product_attention with custom masking.
|
||||
|
||||
Returns:
|
||||
Output tensor of shape [B, L, H*D]
|
||||
"""
|
||||
target_count = q.shape[0] - num_references
|
||||
# Apply rotary positional encoding for Flux models
|
||||
if pe is not None:
|
||||
q, k = flux_math.apply_rope(q, k, pe)
|
||||
|
||||
# Get actual sequence length from QKV tensors
|
||||
actual_seq_len = q.shape[2]
|
||||
|
||||
#if attn_mask is None or attn_mask.shape[-2] != actual_seq_len:
|
||||
# this_source = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1)
|
||||
# this_target = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1)
|
||||
# attn_mask = torch.stack(torch.min(this_source.unsqueeze(1), this_target.unsqueeze(2)), dim=-1).to(device)
|
||||
Args:
|
||||
q: Query tensor of shape [B, H, L, D]
|
||||
k: Key tensor of shape [B, H, L, D]
|
||||
v: Value tensor of shape [B, H, L, D]
|
||||
pe: Positional encoding tensor
|
||||
attn_mask: Optional attention mask of shape [B, H, L, L]
|
||||
|
||||
Returns:
|
||||
Output tensor of shape [B, L, H*D]
|
||||
"""
|
||||
if sampling_step != sampled_step + 1:
|
||||
layer = 0
|
||||
sampling_step = sampled_step + 1
|
||||
if sampling_step >= kv_start_step and sampling_step < kv_end_step:
|
||||
subject[0]["k"] = k
|
||||
subject[0]["v"] = v
|
||||
# Apply rotary positional encoding for Flux models
|
||||
if pe is not None:
|
||||
q, k = flux_math.apply_rope(q, k, pe)
|
||||
|
||||
reference_seq_mask_ = reference_seq_mask.to(device, dtype=attn_dtype)
|
||||
reference_attn_mask = torch.min(reference_seq_mask_.unsqueeze(1), reference_seq_mask_.unsqueeze(2))
|
||||
target_attn_mask = rearrange(torch.min(source_seq_masks.to(device, dtype=attn_dtype).unsqueeze(1), target_seq_masks.to(device, dtype=attn_dtype).unsqueeze(2)), "B H W -> 1 W (B H)")
|
||||
|
||||
x = torch.zeros(q.shape, device=device)
|
||||
# PyTorch's scaled_dot_product_attention
|
||||
x[:num_references] = F.scaled_dot_product_attention(
|
||||
q[:num_references], k[:num_references], v[:num_references],
|
||||
attn_mask=reference_attn_mask
|
||||
)
|
||||
# Get actual sequence length from QKV tensors
|
||||
full_len = q.shape[2]
|
||||
text_len = full_len - h * w # because the latent size is equal to patch count for flux and chroma.
|
||||
|
||||
B, H, L, D = k.shape
|
||||
k = k.permute(1, 0, 2, 3).reshape(1, H, B * L, D)
|
||||
B, H, L, D = v.shape
|
||||
v = v.permute(1, 0, 2, 3).reshape(1, H, B * L, D)
|
||||
|
||||
for i in range(target_count):
|
||||
target_idx = i + num_references
|
||||
if i > 0:
|
||||
k[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = k[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len]
|
||||
v[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = v[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len]
|
||||
|
||||
h[target_idx:target_idx + 1] = F.scaled_dot_product_attention(
|
||||
q[target_idx: target_idx + 1],
|
||||
k[:, :, :num_references*actual_seq_len + actual_seq_len],
|
||||
v[:, :, :num_references*actual_seq_len + actual_seq_len], attn_mask=target_attn_mask)
|
||||
|
||||
# Reshape back to the expected output format
|
||||
h = rearrange(h, "B H L D -> B L (H D)")
|
||||
return h
|
||||
if sampling_step == save_mask_step:
|
||||
subject[0]["front_mask"] = get_fg_attn_weights(text_len, text_len, q, k)
|
||||
|
||||
# Create the attention override function
|
||||
def batched_masked_attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask, transformer_options, reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor:
|
||||
if q.shape[0] != num_references + num_targets or reference_seq_mask is None or reference_seq_mask.shape[-1] != q.shape[2]:
|
||||
return orig_attention(q, k, v, pe)
|
||||
else:
|
||||
h = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=mask
|
||||
)
|
||||
# Reshape back to the expected output format
|
||||
h = rearrange(h, "B H L D -> B L (H D)")
|
||||
layer += 1
|
||||
return h
|
||||
|
||||
# Create the attention override function
|
||||
def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
|
||||
latent_preview.prepare_callback = prepare_callback
|
||||
return masked_attention(q, k, v, pe, reference_seq_mask, target_seq_masks, source_seq_masks)
|
||||
|
||||
# Apply the attention override based on model type
|
||||
overrider = partial(batched_masked_attention_override, reference_seq_mask=reference_seq_mask, target_seq_masks=target_seq_masks, source_seq_masks=source_seq_masks)
|
||||
if model_type == "flux":
|
||||
flux_math.attention = overrider
|
||||
else: # chroma
|
||||
chroma_layers.attention = overrider
|
||||
|
||||
# Clean up
|
||||
#del source_patch_masks, target_patch_masks
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
print('BatchedMaskedReferenceGen processing completed')
|
||||
|
||||
return (model,)
|
||||
return masked_attention(q, k, v, pe, mask)
|
||||
|
||||
# Apply the attention override based on model type
|
||||
overrider = attention_override
|
||||
|
||||
if model.model.__class__.__name__ == "Chroma":
|
||||
model_type = "chroma"
|
||||
|
||||
print(f'GenSubject Detected model type: {model_type}')
|
||||
|
||||
if model_type == "flux":
|
||||
flux_layers.attention = overrider
|
||||
else: # chroma
|
||||
chroma_layers.attention = overrider
|
||||
origin_callback = orig_prepare_callback(model, steps, x0_output_dict)
|
||||
def callback_change_sampled_step(step, x0, x, total_steps):
|
||||
sampled_step = step
|
||||
origin_callback(step, x0, x, total_steps)
|
||||
return callback_change_sampled_step
|
||||
latent_preview.prepare_callback = prepare_callback
|
||||
return (latent,)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user