fix:easy instantIDApplyADV error
This commit is contained in:
yolain
+45
-8
@@ -61,30 +61,67 @@ FOOOCUS_INPAINT_PATCH = {
|
||||
|
||||
LAYER_DIFFUSION_VAE = {
|
||||
"encode": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_encoder.safetensors"
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_encoder.safetensors"
|
||||
}
|
||||
},
|
||||
"decode": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_vae_transparent_decoder.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
LAYER_DIFFUSION = {
|
||||
"Attention Injection": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_attn.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_transparent_attn.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_attn.safetensors"
|
||||
},
|
||||
},
|
||||
"Conv Injection": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors"
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors"
|
||||
},
|
||||
"sd15": {
|
||||
"model_url": None
|
||||
}
|
||||
},
|
||||
"Foreground": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fg2ble.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_joint.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fg2ble.safetensors"
|
||||
}
|
||||
},
|
||||
"Foreground to Background": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fgble2bg.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_fg2bg.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fgble2bg.safetensors"
|
||||
}
|
||||
},
|
||||
"Background": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bg2ble.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_joint.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bg2ble.safetensors"
|
||||
}
|
||||
},
|
||||
"Background to Foreground": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bgble2fg.safetensors"
|
||||
"sd15": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_bg2fg.safetensors"
|
||||
},
|
||||
"sdxl": {
|
||||
"model_url": "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bgble2fg.safetensor"
|
||||
}
|
||||
},
|
||||
}
|
||||
+4
-8
@@ -15,7 +15,7 @@ from .config import MAX_SEED_NUM, BASE_RESOLUTIONS, RESOURCES_DIR, INPAINT_DIR,
|
||||
from .log import log_node_info, log_node_error, log_node_warn
|
||||
from .wildcards import process_with_loras, get_wildcard_list, process
|
||||
from .adv_encode import advanced_encode
|
||||
from .layer_diffusion import LayerDiffuse, LayerMethod, calculate_weight_adjust_channel
|
||||
from .layer_diffuse.func import LayerDiffuse, LayerMethod
|
||||
|
||||
from .libs.utils import find_wildcards_seed, is_linked_styles_selector, easySave, get_local_filepath, add_folder_path_and_extensions
|
||||
from .libs.loader import easyLoader
|
||||
@@ -1687,7 +1687,7 @@ class instantIDApplyAdvanced(instantID):
|
||||
RETURN_NAMES = ("pipe", "model", "positive", "negative")
|
||||
OUTPUT_NODE = True
|
||||
|
||||
FUNCTION = "apply"
|
||||
FUNCTION = "apply_advanced"
|
||||
CATEGORY = "EasyUse/__for_testing"
|
||||
|
||||
def apply_advanced(self, pipe, image, instantid_file, insightface, control_net_name, cn_strength, cn_soft_weights, weight, start_at, end_at, noise, image_kps=None, mask=None, control_net=None, positive=None, negative=None, prompt=None, extra_pnginfo=None, my_unique_id=None):
|
||||
@@ -2524,10 +2524,6 @@ class samplerFull(LayerDiffuse):
|
||||
method = self.get_layer_diffusion_method(pipe['loader_settings']['layer_diffusion_method'], samp_blend_samples is not None)
|
||||
|
||||
weight = pipe['loader_settings']['layer_diffusion_weight'] if 'layer_diffusion_weight' in pipe['loader_settings'] else 1.0
|
||||
try:
|
||||
ModelPatcher.calculate_weight = calculate_weight_adjust_channel(ModelPatcher.calculate_weight)
|
||||
except:
|
||||
pass
|
||||
samp_model, samp_positive, samp_negative = self.apply_layer_diffusion(samp_model, method, weight, samp_samples, samp_blend_samples, samp_positive, samp_negative)
|
||||
|
||||
def downscale_model_unet(samp_model):
|
||||
@@ -2594,7 +2590,7 @@ class samplerFull(LayerDiffuse):
|
||||
samp_images = samp_vae.decode(latent).cpu()
|
||||
|
||||
# LayerDiffusion Decode
|
||||
new_images, samp_images, alpha = self.layer_diffusion_decode(layer_diffusion_method, latent, blend_samples, samp_images)
|
||||
new_images, samp_images, alpha = self.layer_diffusion_decode(layer_diffusion_method, latent, blend_samples, samp_images, samp_model)
|
||||
|
||||
# 推理总耗时（包含解码）
|
||||
end_decode_time = int(time.time() * 1000)
|
||||
@@ -2715,7 +2711,7 @@ class samplerFull(LayerDiffuse):
|
||||
output_images = torch.stack([tensor.squeeze() for tensor in image_list])
|
||||
|
||||
new_images, samp_images, alpha = self.layer_diffusion_decode(layer_diffusion_method, latents_plot, blend_samples,
|
||||
output_images)
|
||||
output_images, samp_model)
|
||||
|
||||
results = easySave(images, save_prefix, image_output, prompt, extra_pnginfo)
|
||||
sampler.update_value_by_id("results", my_unique_id, results)
|
||||
|
||||
@@ -0,0 +1,360 @@
|
||||
# Currently only sd15
|
||||
|
||||
import functools
|
||||
import torch
|
||||
import einops
|
||||
|
||||
from comfy import model_management, utils
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
|
||||
|
||||
module_mapping_sd15 = {
|
||||
0: "input_blocks.1.1.transformer_blocks.0.attn1",
|
||||
1: "input_blocks.1.1.transformer_blocks.0.attn2",
|
||||
2: "input_blocks.2.1.transformer_blocks.0.attn1",
|
||||
3: "input_blocks.2.1.transformer_blocks.0.attn2",
|
||||
4: "input_blocks.4.1.transformer_blocks.0.attn1",
|
||||
5: "input_blocks.4.1.transformer_blocks.0.attn2",
|
||||
6: "input_blocks.5.1.transformer_blocks.0.attn1",
|
||||
7: "input_blocks.5.1.transformer_blocks.0.attn2",
|
||||
8: "input_blocks.7.1.transformer_blocks.0.attn1",
|
||||
9: "input_blocks.7.1.transformer_blocks.0.attn2",
|
||||
10: "input_blocks.8.1.transformer_blocks.0.attn1",
|
||||
11: "input_blocks.8.1.transformer_blocks.0.attn2",
|
||||
12: "output_blocks.3.1.transformer_blocks.0.attn1",
|
||||
13: "output_blocks.3.1.transformer_blocks.0.attn2",
|
||||
14: "output_blocks.4.1.transformer_blocks.0.attn1",
|
||||
15: "output_blocks.4.1.transformer_blocks.0.attn2",
|
||||
16: "output_blocks.5.1.transformer_blocks.0.attn1",
|
||||
17: "output_blocks.5.1.transformer_blocks.0.attn2",
|
||||
18: "output_blocks.6.1.transformer_blocks.0.attn1",
|
||||
19: "output_blocks.6.1.transformer_blocks.0.attn2",
|
||||
20: "output_blocks.7.1.transformer_blocks.0.attn1",
|
||||
21: "output_blocks.7.1.transformer_blocks.0.attn2",
|
||||
22: "output_blocks.8.1.transformer_blocks.0.attn1",
|
||||
23: "output_blocks.8.1.transformer_blocks.0.attn2",
|
||||
24: "output_blocks.9.1.transformer_blocks.0.attn1",
|
||||
25: "output_blocks.9.1.transformer_blocks.0.attn2",
|
||||
26: "output_blocks.10.1.transformer_blocks.0.attn1",
|
||||
27: "output_blocks.10.1.transformer_blocks.0.attn2",
|
||||
28: "output_blocks.11.1.transformer_blocks.0.attn1",
|
||||
29: "output_blocks.11.1.transformer_blocks.0.attn2",
|
||||
30: "middle_block.1.transformer_blocks.0.attn1",
|
||||
31: "middle_block.1.transformer_blocks.0.attn2",
|
||||
}
|
||||
|
||||
|
||||
def compute_cond_mark(cond_or_uncond, sigmas):
|
||||
cond_or_uncond_size = int(sigmas.shape[0])
|
||||
|
||||
cond_mark = []
|
||||
for cx in cond_or_uncond:
|
||||
cond_mark += [cx] * cond_or_uncond_size
|
||||
|
||||
cond_mark = torch.Tensor(cond_mark).to(sigmas)
|
||||
return cond_mark
|
||||
|
||||
|
||||
class LoRALinearLayer(torch.nn.Module):
|
||||
def __init__(self, in_features: int, out_features: int, rank: int = 256, org=None):
|
||||
super().__init__()
|
||||
self.down = torch.nn.Linear(in_features, rank, bias=False)
|
||||
self.up = torch.nn.Linear(rank, out_features, bias=False)
|
||||
self.org = [org]
|
||||
|
||||
def forward(self, h):
|
||||
org_weight = self.org[0].weight.to(h)
|
||||
org_bias = self.org[0].bias.to(h) if self.org[0].bias is not None else None
|
||||
down_weight = self.down.weight
|
||||
up_weight = self.up.weight
|
||||
final_weight = org_weight + torch.mm(up_weight, down_weight)
|
||||
return torch.nn.functional.linear(h, final_weight, org_bias)
|
||||
|
||||
|
||||
class AttentionSharingUnit(torch.nn.Module):
|
||||
# `transformer_options` passed to the most recent BasicTransformerBlock.forward
|
||||
# call.
|
||||
transformer_options: dict = {}
|
||||
|
||||
def __init__(self, module, frames=2, use_control=True, rank=256):
|
||||
super().__init__()
|
||||
|
||||
self.heads = module.heads
|
||||
self.frames = frames
|
||||
self.original_module = [module]
|
||||
q_in_channels, q_out_channels = (
|
||||
module.to_q.in_features,
|
||||
module.to_q.out_features,
|
||||
)
|
||||
k_in_channels, k_out_channels = (
|
||||
module.to_k.in_features,
|
||||
module.to_k.out_features,
|
||||
)
|
||||
v_in_channels, v_out_channels = (
|
||||
module.to_v.in_features,
|
||||
module.to_v.out_features,
|
||||
)
|
||||
o_in_channels, o_out_channels = (
|
||||
module.to_out[0].in_features,
|
||||
module.to_out[0].out_features,
|
||||
)
|
||||
|
||||
hidden_size = k_out_channels
|
||||
|
||||
self.to_q_lora = [
|
||||
LoRALinearLayer(q_in_channels, q_out_channels, rank, module.to_q)
|
||||
for _ in range(self.frames)
|
||||
]
|
||||
self.to_k_lora = [
|
||||
LoRALinearLayer(k_in_channels, k_out_channels, rank, module.to_k)
|
||||
for _ in range(self.frames)
|
||||
]
|
||||
self.to_v_lora = [
|
||||
LoRALinearLayer(v_in_channels, v_out_channels, rank, module.to_v)
|
||||
for _ in range(self.frames)
|
||||
]
|
||||
self.to_out_lora = [
|
||||
LoRALinearLayer(o_in_channels, o_out_channels, rank, module.to_out[0])
|
||||
for _ in range(self.frames)
|
||||
]
|
||||
|
||||
self.to_q_lora = torch.nn.ModuleList(self.to_q_lora)
|
||||
self.to_k_lora = torch.nn.ModuleList(self.to_k_lora)
|
||||
self.to_v_lora = torch.nn.ModuleList(self.to_v_lora)
|
||||
self.to_out_lora = torch.nn.ModuleList(self.to_out_lora)
|
||||
|
||||
self.temporal_i = torch.nn.Linear(
|
||||
in_features=hidden_size, out_features=hidden_size
|
||||
)
|
||||
self.temporal_n = torch.nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6
|
||||
)
|
||||
self.temporal_q = torch.nn.Linear(
|
||||
in_features=hidden_size, out_features=hidden_size
|
||||
)
|
||||
self.temporal_k = torch.nn.Linear(
|
||||
in_features=hidden_size, out_features=hidden_size
|
||||
)
|
||||
self.temporal_v = torch.nn.Linear(
|
||||
in_features=hidden_size, out_features=hidden_size
|
||||
)
|
||||
self.temporal_o = torch.nn.Linear(
|
||||
in_features=hidden_size, out_features=hidden_size
|
||||
)
|
||||
|
||||
self.control_convs = None
|
||||
|
||||
if use_control:
|
||||
self.control_convs = [
|
||||
torch.nn.Sequential(
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(256, hidden_size, kernel_size=1),
|
||||
)
|
||||
for _ in range(self.frames)
|
||||
]
|
||||
self.control_convs = torch.nn.ModuleList(self.control_convs)
|
||||
|
||||
self.control_signals = None
|
||||
|
||||
def forward(self, h, context=None, value=None):
|
||||
transformer_options = self.transformer_options
|
||||
|
||||
modified_hidden_states = einops.rearrange(
|
||||
h, "(b f) d c -> f b d c", f=self.frames
|
||||
)
|
||||
|
||||
if self.control_convs is not None:
|
||||
context_dim = int(modified_hidden_states.shape[2])
|
||||
control_outs = []
|
||||
for f in range(self.frames):
|
||||
control_signal = self.control_signals[context_dim].to(
|
||||
modified_hidden_states
|
||||
)
|
||||
control = self.control_convs[f](control_signal)
|
||||
control = einops.rearrange(control, "b c h w -> b (h w) c")
|
||||
control_outs.append(control)
|
||||
control_outs = torch.stack(control_outs, dim=0)
|
||||
modified_hidden_states = modified_hidden_states + control_outs.to(
|
||||
modified_hidden_states
|
||||
)
|
||||
|
||||
if context is None:
|
||||
framed_context = modified_hidden_states
|
||||
else:
|
||||
framed_context = einops.rearrange(
|
||||
context, "(b f) d c -> f b d c", f=self.frames
|
||||
)
|
||||
|
||||
framed_cond_mark = einops.rearrange(
|
||||
compute_cond_mark(
|
||||
transformer_options["cond_or_uncond"],
|
||||
transformer_options["sigmas"],
|
||||
),
|
||||
"(b f) -> f b",
|
||||
f=self.frames,
|
||||
).to(modified_hidden_states)
|
||||
|
||||
attn_outs = []
|
||||
for f in range(self.frames):
|
||||
fcf = framed_context[f]
|
||||
|
||||
if context is not None:
|
||||
cond_overwrite = transformer_options.get("cond_overwrite", [])
|
||||
if len(cond_overwrite) > f:
|
||||
cond_overwrite = cond_overwrite[f]
|
||||
else:
|
||||
cond_overwrite = None
|
||||
if cond_overwrite is not None:
|
||||
cond_mark = framed_cond_mark[f][:, None, None]
|
||||
fcf = cond_overwrite.to(fcf) * (1.0 - cond_mark) + fcf * cond_mark
|
||||
|
||||
q = self.to_q_lora[f](modified_hidden_states[f])
|
||||
k = self.to_k_lora[f](fcf)
|
||||
v = self.to_v_lora[f](fcf)
|
||||
o = optimized_attention(q, k, v, self.heads)
|
||||
o = self.to_out_lora[f](o)
|
||||
o = self.original_module[0].to_out[1](o)
|
||||
attn_outs.append(o)
|
||||
|
||||
attn_outs = torch.stack(attn_outs, dim=0)
|
||||
modified_hidden_states = modified_hidden_states + attn_outs.to(
|
||||
modified_hidden_states
|
||||
)
|
||||
modified_hidden_states = einops.rearrange(
|
||||
modified_hidden_states, "f b d c -> (b f) d c", f=self.frames
|
||||
)
|
||||
|
||||
x = modified_hidden_states
|
||||
x = self.temporal_n(x)
|
||||
x = self.temporal_i(x)
|
||||
d = x.shape[1]
|
||||
|
||||
x = einops.rearrange(x, "(b f) d c -> (b d) f c", f=self.frames)
|
||||
|
||||
q = self.temporal_q(x)
|
||||
k = self.temporal_k(x)
|
||||
v = self.temporal_v(x)
|
||||
|
||||
x = optimized_attention(q, k, v, self.heads)
|
||||
x = self.temporal_o(x)
|
||||
x = einops.rearrange(x, "(b d) f c -> (b f) d c", d=d)
|
||||
|
||||
modified_hidden_states = modified_hidden_states + x
|
||||
|
||||
return modified_hidden_states - h
|
||||
|
||||
@classmethod
|
||||
def hijack_transformer_block(cls):
|
||||
def register_get_transformer_options(func):
|
||||
@functools.wraps(func)
|
||||
def forward(self, x, context=None, transformer_options={}):
|
||||
cls.transformer_options = transformer_options
|
||||
return func(self, x, context, transformer_options)
|
||||
|
||||
return forward
|
||||
|
||||
from comfy.ldm.modules.attention import BasicTransformerBlock
|
||||
|
||||
BasicTransformerBlock.forward = register_get_transformer_options(
|
||||
BasicTransformerBlock.forward
|
||||
)
|
||||
|
||||
|
||||
AttentionSharingUnit.hijack_transformer_block()
|
||||
|
||||
|
||||
class AdditionalAttentionCondsEncoder(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.blocks_0 = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(3, 32, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
) # 64*64*256
|
||||
|
||||
self.blocks_1 = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
) # 32*32*256
|
||||
|
||||
self.blocks_2 = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
) # 16*16*256
|
||||
|
||||
self.blocks_3 = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
|
||||
torch.nn.SiLU(),
|
||||
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
|
||||
torch.nn.SiLU(),
|
||||
) # 8*8*256
|
||||
|
||||
self.blks = [self.blocks_0, self.blocks_1, self.blocks_2, self.blocks_3]
|
||||
|
||||
def __call__(self, h):
|
||||
results = {}
|
||||
for b in self.blks:
|
||||
h = b(h)
|
||||
results[int(h.shape[2]) * int(h.shape[3])] = h
|
||||
return results
|
||||
|
||||
|
||||
class HookerLayers(torch.nn.Module):
|
||||
def __init__(self, layer_list):
|
||||
super().__init__()
|
||||
self.layers = torch.nn.ModuleList(layer_list)
|
||||
|
||||
|
||||
class AttentionSharingPatcher(torch.nn.Module):
|
||||
def __init__(self, unet, frames=2, use_control=True, rank=256):
|
||||
super().__init__()
|
||||
model_management.unload_model_clones(unet)
|
||||
|
||||
units = []
|
||||
for i in range(32):
|
||||
real_key = module_mapping_sd15[i]
|
||||
attn_module = utils.get_attr(unet.model.diffusion_model, real_key)
|
||||
u = AttentionSharingUnit(
|
||||
attn_module, frames=frames, use_control=use_control, rank=rank
|
||||
)
|
||||
units.append(u)
|
||||
unet.add_object_patch("diffusion_model." + real_key, u)
|
||||
|
||||
self.hookers = HookerLayers(units)
|
||||
|
||||
if use_control:
|
||||
self.kwargs_encoder = AdditionalAttentionCondsEncoder()
|
||||
else:
|
||||
self.kwargs_encoder = None
|
||||
|
||||
self.dtype = torch.float32
|
||||
if model_management.should_use_fp16(model_management.get_torch_device()):
|
||||
self.dtype = torch.float16
|
||||
self.hookers.half()
|
||||
return
|
||||
|
||||
def set_control(self, img):
|
||||
img = img.cpu().float() * 2.0 - 1.0
|
||||
signals = self.kwargs_encoder(img)
|
||||
for m in self.hookers.layers:
|
||||
m.control_signals = signals
|
||||
return
|
||||
@@ -0,0 +1,138 @@
|
||||
import torch
|
||||
import comfy.model_management
|
||||
from enum import Enum
|
||||
from comfy.utils import load_torch_file
|
||||
from comfy.conds import CONDRegular
|
||||
from comfy_extras.nodes_compositing import JoinImageWithAlpha
|
||||
from .model import ModelPatcher, TransparentVAEDecoder, calculate_weight_adjust_channel
|
||||
from .attension_sharing import AttentionSharingPatcher
|
||||
from ..config import LAYER_DIFFUSION, LAYER_DIFFUSION_DIR, LAYER_DIFFUSION_VAE
|
||||
from ..libs.utils import to_lora_patch_dict, get_local_filepath, get_sd_version
|
||||
|
||||
class LayerMethod(Enum):
|
||||
FG_ONLY_ATTN = "Attention Injection"
|
||||
FG_ONLY_CONV = "Conv Injection"
|
||||
FG_TO_BLEND = "Foreground"
|
||||
FG_BLEND_TO_BG = "Foreground to Background"
|
||||
BG_TO_BLEND = "Background"
|
||||
BG_BLEND_TO_FG = "Background to Foreground"
|
||||
|
||||
class LayerDiffuse:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.vae_transparent_decoder = None
|
||||
self.frames = 1
|
||||
try:
|
||||
ModelPatcher.calculate_weight = calculate_weight_adjust_channel(ModelPatcher.calculate_weight)
|
||||
except:
|
||||
pass
|
||||
|
||||
def get_layer_diffusion_method(self, method, has_blend_latent):
|
||||
method = LayerMethod(method)
|
||||
if method == LayerMethod.BG_TO_BLEND and has_blend_latent:
|
||||
method = LayerMethod.BG_BLEND_TO_FG
|
||||
elif method == LayerMethod.FG_TO_BLEND and has_blend_latent:
|
||||
method = LayerMethod.FG_BLEND_TO_BG
|
||||
return method
|
||||
|
||||
def apply_layer_c_concat(self, cond, uncond, c_concat):
|
||||
def write_c_concat(cond):
|
||||
new_cond = []
|
||||
for t in cond:
|
||||
n = [t[0], t[1].copy()]
|
||||
if "model_conds" not in n[1]:
|
||||
n[1]["model_conds"] = {}
|
||||
n[1]["model_conds"]["c_concat"] = CONDRegular(c_concat)
|
||||
new_cond.append(n)
|
||||
return new_cond
|
||||
|
||||
return (write_c_concat(cond), write_c_concat(uncond))
|
||||
|
||||
def apply_layer_diffusion(self, model: ModelPatcher, method, weight, samples, blend_samples, positive, negative, control_img=None):
|
||||
sd_version = get_sd_version(model)
|
||||
model_url = LAYER_DIFFUSION[method.value][sd_version]["model_url"]
|
||||
if method in [LayerMethod.FG_ONLY_CONV, LayerMethod.FG_ONLY_ATTN] and sd_version == 'sd15':
|
||||
self.frames = 3
|
||||
if method == LayerMethod.BG_BLEND_TO_FG and sd_version == 'sd15':
|
||||
self.frames = 2
|
||||
if model_url is None:
|
||||
raise Exception(f"{method.value} is not supported for {sd_version} model")
|
||||
model_file = get_local_filepath(model_url, LAYER_DIFFUSION_DIR)
|
||||
|
||||
layer_lora_state_dict = load_torch_file(model_file)
|
||||
work_model = model.clone()
|
||||
if sd_version == 'sd15':
|
||||
patcher = AttentionSharingPatcher(
|
||||
work_model, self.frames, use_control=control_img is not None
|
||||
)
|
||||
patcher.load_state_dict(layer_lora_state_dict, strict=True)
|
||||
if control_img is not None:
|
||||
patcher.set_control(control_img)
|
||||
else:
|
||||
layer_lora_patch_dict = to_lora_patch_dict(layer_lora_state_dict)
|
||||
work_model.add_patches(layer_lora_patch_dict, weight)
|
||||
|
||||
# cond_contact
|
||||
if method in [LayerMethod.FG_ONLY_ATTN, LayerMethod.FG_ONLY_CONV]:
|
||||
samp_model = work_model
|
||||
else:
|
||||
if method in [LayerMethod.BG_TO_BLEND, LayerMethod.FG_TO_BLEND]:
|
||||
c_concat = model.model.latent_format.process_in(samples["samples"])
|
||||
else:
|
||||
c_concat = model.model.latent_format.process_in(torch.cat([samples["samples"], blend_samples["samples"]], dim=1))
|
||||
samp_model, positive, negative = (work_model,) + self.apply_layer_c_concat(positive, negative, c_concat)
|
||||
|
||||
return samp_model, positive, negative
|
||||
|
||||
def join_image_with_alpha(self, image, alpha):
|
||||
out = image.movedim(-1, 1)
|
||||
if out.shape[1] == 3: # RGB
|
||||
out = torch.cat([out, torch.ones_like(out[:, :1, :, :])], dim=1)
|
||||
for i in range(out.shape[0]):
|
||||
out[i, 3, :, :] = alpha
|
||||
return out.movedim(1, -1)
|
||||
|
||||
def layer_diffusion_decode(self, layer_diffusion_method, latent, blend_samples, samp_images, model):
|
||||
alpha = None
|
||||
if layer_diffusion_method is not None:
|
||||
method = self.get_layer_diffusion_method(layer_diffusion_method, blend_samples is not None)
|
||||
print(method.value)
|
||||
if method in [LayerMethod.FG_ONLY_CONV, LayerMethod.FG_ONLY_ATTN, LayerMethod.BG_BLEND_TO_FG]:
|
||||
if self.vae_transparent_decoder is None:
|
||||
sd_version = get_sd_version(model)
|
||||
print(sd_version)
|
||||
if sd_version not in ['sdxl', 'sd15']:
|
||||
raise Exception(f"Only SDXL and SD1.5 model supported for Layer Diffusion")
|
||||
model_url = LAYER_DIFFUSION_VAE['decode'][sd_version]["model_url"]
|
||||
if model_url is None:
|
||||
raise Exception(f"{method.value} is not supported for {sd_version} model")
|
||||
decoder_file = get_local_filepath(model_url, LAYER_DIFFUSION_DIR)
|
||||
self.vae_transparent_decoder = TransparentVAEDecoder(
|
||||
load_torch_file(decoder_file),
|
||||
device=comfy.model_management.get_torch_device(),
|
||||
dtype=(torch.float16 if comfy.model_management.should_use_fp16() else torch.float32),
|
||||
)
|
||||
|
||||
pixel = samp_images.movedim(-1, 1) # [B, H, W, C] => [B, C, H, W]
|
||||
decoded = []
|
||||
sub_batch_size = 16
|
||||
for start_idx in range(0, latent.shape[0], sub_batch_size):
|
||||
decoded.append(
|
||||
self.vae_transparent_decoder.decode_pixel(
|
||||
pixel[start_idx: start_idx + sub_batch_size],
|
||||
latent[start_idx: start_idx + sub_batch_size],
|
||||
)
|
||||
)
|
||||
pixel_with_alpha = torch.cat(decoded, dim=0)
|
||||
# [B, C, H, W] => [B, H, W, C]
|
||||
pixel_with_alpha = pixel_with_alpha.movedim(1, -1)
|
||||
image = pixel_with_alpha[..., 1:]
|
||||
alpha = pixel_with_alpha[..., 0]
|
||||
|
||||
alpha = 1.0 - alpha
|
||||
new_images, = JoinImageWithAlpha().join_image_with_alpha(image, alpha)
|
||||
else:
|
||||
new_images = samp_images
|
||||
else:
|
||||
new_images = samp_images
|
||||
return (new_images, samp_images, alpha)
|
||||
@@ -5,17 +5,9 @@ import numpy as np
|
||||
import comfy.model_management
|
||||
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from enum import Enum
|
||||
from tqdm import tqdm
|
||||
from typing import Optional, Tuple
|
||||
|
||||
class LayerMethod(Enum):
|
||||
FG_ONLY_ATTN = "Attention Injection"
|
||||
FG_ONLY_CONV = "Conv Injection"
|
||||
FG_TO_BLEND = "Foreground"
|
||||
FG_BLEND_TO_BG = "Foreground to Background"
|
||||
BG_TO_BLEND = "Background"
|
||||
BG_BLEND_TO_FG = "Background to Foreground"
|
||||
|
||||
try:
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
@@ -383,100 +375,4 @@ except ImportError:
|
||||
print("\33[31mpip install diffusers\033[0m")
|
||||
|
||||
|
||||
from comfy.utils import load_torch_file
|
||||
from comfy.conds import CONDRegular
|
||||
from comfy_extras.nodes_compositing import JoinImageWithAlpha
|
||||
from .config import LAYER_DIFFUSION, LAYER_DIFFUSION_DIR, LAYER_DIFFUSION_VAE
|
||||
from .libs.utils import to_lora_patch_dict, get_local_filepath
|
||||
class LayerDiffuse:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.vae_transparent_decoder = None
|
||||
self.vae_transparent_encoder = None
|
||||
|
||||
def get_layer_diffusion_method(self, method, has_blend_latent):
|
||||
method = LayerMethod(method)
|
||||
if method == LayerMethod.BG_TO_BLEND and has_blend_latent:
|
||||
method = LayerMethod.BG_BLEND_TO_FG
|
||||
elif method == LayerMethod.FG_TO_BLEND and has_blend_latent:
|
||||
method = LayerMethod.FG_BLEND_TO_BG
|
||||
return method
|
||||
|
||||
def apply_layer_c_concat(self, cond, uncond, c_concat):
|
||||
def write_c_concat(cond):
|
||||
new_cond = []
|
||||
for t in cond:
|
||||
n = [t[0], t[1].copy()]
|
||||
if "model_conds" not in n[1]:
|
||||
n[1]["model_conds"] = {}
|
||||
n[1]["model_conds"]["c_concat"] = CONDRegular(c_concat)
|
||||
new_cond.append(n)
|
||||
return new_cond
|
||||
|
||||
return (write_c_concat(cond), write_c_concat(uncond))
|
||||
|
||||
def apply_layer_diffusion(self, model: ModelPatcher, method, weight, samples, blend_samples, positive, negative):
|
||||
|
||||
model_file = get_local_filepath(LAYER_DIFFUSION[method.value]["model_url"], LAYER_DIFFUSION_DIR)
|
||||
layer_lora_state_dict = load_torch_file(model_file)
|
||||
layer_lora_patch_dict = to_lora_patch_dict(layer_lora_state_dict)
|
||||
work_model = model.clone()
|
||||
work_model.add_patches(layer_lora_patch_dict, weight)
|
||||
|
||||
# cond_contact
|
||||
if method in [LayerMethod.FG_ONLY_ATTN, LayerMethod.FG_ONLY_CONV]:
|
||||
samp_model = work_model
|
||||
else:
|
||||
if method in [LayerMethod.BG_TO_BLEND, LayerMethod.FG_TO_BLEND]:
|
||||
c_concat = model.model.latent_format.process_in(samples["samples"])
|
||||
else:
|
||||
c_concat = model.model.latent_format.process_in(torch.cat([samples["samples"], blend_samples["samples"]], dim=1))
|
||||
samp_model, positive, negative = (work_model,) + self.apply_layer_c_concat(positive, negative, c_concat)
|
||||
|
||||
return samp_model, positive, negative
|
||||
|
||||
def join_image_with_alpha(self, image, alpha):
|
||||
out = image.movedim(-1, 1)
|
||||
if out.shape[1] == 3: # RGB
|
||||
out = torch.cat([out, torch.ones_like(out[:, :1, :, :])], dim=1)
|
||||
for i in range(out.shape[0]):
|
||||
out[i, 3, :, :] = alpha
|
||||
return out.movedim(1, -1)
|
||||
|
||||
def layer_diffusion_decode(self, layer_diffusion_method, latent, blend_samples, samp_images):
|
||||
alpha = None
|
||||
if layer_diffusion_method is not None:
|
||||
method = self.get_layer_diffusion_method(layer_diffusion_method, blend_samples is not None)
|
||||
print(method.value)
|
||||
if method in [LayerMethod.FG_ONLY_CONV, LayerMethod.FG_ONLY_ATTN, LayerMethod.BG_BLEND_TO_FG]:
|
||||
if self.vae_transparent_decoder is None:
|
||||
decoder_file = get_local_filepath(LAYER_DIFFUSION_VAE['decode']["model_url"], LAYER_DIFFUSION_DIR)
|
||||
self.vae_transparent_decoder = TransparentVAEDecoder(
|
||||
load_torch_file(decoder_file),
|
||||
device=comfy.model_management.get_torch_device(),
|
||||
dtype=(torch.float16 if comfy.model_management.should_use_fp16() else torch.float32),
|
||||
)
|
||||
|
||||
pixel = samp_images.movedim(-1, 1) # [B, H, W, C] => [B, C, H, W]
|
||||
decoded = []
|
||||
sub_batch_size = 16
|
||||
for start_idx in range(0, latent.shape[0], sub_batch_size):
|
||||
decoded.append(
|
||||
self.vae_transparent_decoder.decode_pixel(
|
||||
pixel[start_idx: start_idx + sub_batch_size],
|
||||
latent[start_idx: start_idx + sub_batch_size],
|
||||
)
|
||||
)
|
||||
pixel_with_alpha = torch.cat(decoded, dim=0)
|
||||
# [B, C, H, W] => [B, H, W, C]
|
||||
pixel_with_alpha = pixel_with_alpha.movedim(1, -1)
|
||||
image = pixel_with_alpha[..., 1:]
|
||||
alpha = pixel_with_alpha[..., 0]
|
||||
|
||||
alpha = 1.0 - alpha
|
||||
new_images, = JoinImageWithAlpha().join_image_with_alpha(image, alpha)
|
||||
else:
|
||||
new_images = samp_images
|
||||
else:
|
||||
new_images = samp_images
|
||||
return (new_images, samp_images, alpha)
|
||||
@@ -27,6 +27,21 @@ def add_folder_path_and_extensions(folder_name, full_folder_paths, extensions):
|
||||
else:
|
||||
folder_paths.folder_names_and_paths[folder_name] = (full_folder_paths, extensions)
|
||||
|
||||
from comfy.model_base import BaseModel
|
||||
import comfy.supported_models
|
||||
import comfy.supported_models_base
|
||||
def get_sd_version(model):
|
||||
base: BaseModel = model.model
|
||||
model_config: comfy.supported_models.supported_models_base.BASE = base.model_config
|
||||
if isinstance(model_config, comfy.supported_models.SDXL):
|
||||
return 'sdxl'
|
||||
elif isinstance(
|
||||
model_config, (comfy.supported_models.SD15, comfy.supported_models.SD20)
|
||||
):
|
||||
return 'sd15'
|
||||
else:
|
||||
return 'unknown'
|
||||
|
||||
def find_nearest_steps(clip_id, prompt):
|
||||
"""Find the nearest KSampler or preSampling node that references the given id."""
|
||||
def check_link_to_clip(node_id, clip_id, visited=None, node=None):
|
||||
|
||||
+1
-1
@@ -1,2 +1,2 @@
|
||||
diffusers==0.25.0
|
||||
diffusers>=0.25.0
|
||||
aiohttp
|
||||
Reference in New Issue
Block a user