275 lines
11 KiB
Python
275 lines
11 KiB
Python
import torch
|
|
import math
|
|
import torch.nn.functional as F
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class YC_LG_Redux:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"conditioning": ("CONDITIONING", ),
|
|
"style_model": ("STYLE_MODEL", ),
|
|
"clip_vision": ("CLIP_VISION",),
|
|
"image": ("IMAGE",),
|
|
"crop": (["center", "mask_area", "none"], {
|
|
"default": "none",
|
|
"tooltip": "裁剪模式：center-中心裁剪, mask_area-遮罩区域裁剪, none-不裁剪"
|
|
}),
|
|
"sharpen": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -5.0,
|
|
"max": 5.0,
|
|
"step": 0.1,
|
|
"tooltip": "锐化强度：负值为模糊，正值为锐化，0为不处理"
|
|
}),
|
|
"patch_res": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1,
|
|
"tooltip": "patch分辨率，数值越大分块越细致"
|
|
}),
|
|
"style_strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.01,
|
|
"tooltip": "风格强度，越高越偏向参考图片"
|
|
}),
|
|
"prompt_strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.01,
|
|
"tooltip": "文本提示词强度，越高文本特征越强"
|
|
}),
|
|
"blend_mode": (["lerp", "feature_boost", "frequency"], {
|
|
"default": "lerp",
|
|
"tooltip": "风格强度的计算方式：\n" +
|
|
"lerp - 线性混合 - 高度参考原图\n" +
|
|
"feature_boost - 特征增强 - 增强真实感\n" +
|
|
"frequency - 频率增强 - 增强高频细节"
|
|
}),
|
|
"noise_level": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "添加随机噪声的强度，可用于修复错误细节"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", ),
|
|
}}
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "apply_stylemodel"
|
|
CATEGORY = "conditioning/style_model"
|
|
|
|
def crop_to_mask_area(self, image, mask):
|
|
if len(image.shape) == 4:
|
|
B, H, W, C = image.shape
|
|
image = image.squeeze(0)
|
|
else:
|
|
H, W, C = image.shape
|
|
|
|
if len(mask.shape) == 3:
|
|
mask = mask.squeeze(0)
|
|
|
|
nonzero_coords = torch.nonzero(mask)
|
|
if len(nonzero_coords) == 0:
|
|
return image, mask
|
|
|
|
top = nonzero_coords[:, 0].min().item()
|
|
bottom = nonzero_coords[:, 0].max().item()
|
|
left = nonzero_coords[:, 1].min().item()
|
|
right = nonzero_coords[:, 1].max().item()
|
|
|
|
width = right - left
|
|
height = bottom - top
|
|
size = max(width, height)
|
|
|
|
center_y = (top + bottom) // 2
|
|
center_x = (left + right) // 2
|
|
|
|
half_size = size // 2
|
|
new_top = max(0, center_y - half_size)
|
|
new_bottom = min(H, center_y + half_size)
|
|
new_left = max(0, center_x - half_size)
|
|
new_right = min(W, center_x + half_size)
|
|
|
|
cropped_image = image[new_top:new_bottom, new_left:new_right]
|
|
cropped_mask = mask[new_top:new_bottom, new_left:new_right]
|
|
|
|
cropped_image = cropped_image.unsqueeze(0)
|
|
cropped_mask = cropped_mask.unsqueeze(0)
|
|
|
|
return cropped_image, cropped_mask
|
|
|
|
def apply_image_preprocess(self, image, strength):
|
|
original_shape = image.shape
|
|
original_device = image.device
|
|
|
|
if torch.is_tensor(image):
|
|
if len(image.shape) == 4:
|
|
image_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
else:
|
|
image_np = (image.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
if strength < 0:
|
|
abs_strength = abs(strength)
|
|
kernel_size = int(3 + abs_strength * 12) // 2 * 2 + 1
|
|
sigma = 0.3 + abs_strength * 2.7
|
|
processed = cv2.GaussianBlur(image_np, (kernel_size, kernel_size), sigma)
|
|
elif strength > 0:
|
|
kernel = np.array([[-1,-1,-1],
|
|
[-1, 9,-1],
|
|
[-1,-1,-1]]) * strength + np.array([[0,0,0],
|
|
[0,1,0],
|
|
[0,0,0]]) * (1 - strength)
|
|
processed = cv2.filter2D(image_np, -1, kernel)
|
|
processed = np.clip(processed, 0, 255)
|
|
else:
|
|
processed = image_np
|
|
|
|
processed_tensor = torch.from_numpy(processed.astype(np.float32) / 255.0).to(original_device)
|
|
if len(original_shape) == 4:
|
|
processed_tensor = processed_tensor.unsqueeze(0)
|
|
|
|
return processed_tensor
|
|
|
|
def apply_style_strength(self, cond, txt, strength, mode="lerp"):
|
|
if mode == "lerp":
|
|
if txt.shape[1] != cond.shape[1]:
|
|
txt_mean = txt.mean(dim=1, keepdim=True)
|
|
txt_expanded = txt_mean.expand(-1, cond.shape[1], -1)
|
|
return torch.lerp(txt_expanded, cond, strength)
|
|
return torch.lerp(txt, cond, strength)
|
|
|
|
elif mode == "feature_boost":
|
|
mean = torch.mean(cond, dim=-1, keepdim=True)
|
|
std = torch.std(cond, dim=-1, keepdim=True)
|
|
normalized = (cond - mean) / (std + 1e-6)
|
|
boost = torch.tanh(normalized * (strength * 2.0))
|
|
return cond * (1 + boost * 2.0)
|
|
|
|
elif mode == "frequency":
|
|
try:
|
|
B, N, C = cond.shape
|
|
x = cond.float()
|
|
fft = torch.fft.rfft(x, dim=-1)
|
|
magnitudes = torch.abs(fft)
|
|
phases = torch.angle(fft)
|
|
freq_dim = fft.shape[-1]
|
|
freq_range = torch.linspace(0, 1, freq_dim, device=cond.device)
|
|
alpha = 2.0 * strength
|
|
beta = 0.5
|
|
filter_response = 1.0 + alpha * torch.pow(freq_range, beta)
|
|
filter_response = filter_response.view(1, 1, -1)
|
|
enhanced_magnitudes = magnitudes * filter_response
|
|
enhanced_fft = enhanced_magnitudes * torch.exp(1j * phases)
|
|
enhanced = torch.fft.irfft(enhanced_fft, n=C, dim=-1)
|
|
mean = enhanced.mean(dim=-1, keepdim=True)
|
|
std = enhanced.std(dim=-1, keepdim=True)
|
|
enhanced_norm = (enhanced - mean) / (std + 1e-6)
|
|
mix_ratio = torch.sigmoid(torch.tensor(strength * 2 - 1))
|
|
result = torch.lerp(cond, enhanced_norm.to(cond.dtype), mix_ratio)
|
|
residual = (result - cond) * strength
|
|
final = cond + residual
|
|
return final
|
|
except Exception as e:
|
|
print(f"频率处理出错: {e}")
|
|
print(f"输入张量形状: {cond.shape}")
|
|
return cond
|
|
|
|
return cond
|
|
|
|
def apply_stylemodel(self, conditioning, style_model, clip_vision, image,
|
|
patch_res=16, style_strength=1.0, prompt_strength=1.0,
|
|
noise_level=0.0, crop="none", sharpen=0.0,
|
|
blend_mode="lerp", mask=None):
|
|
|
|
processed_image = image.clone()
|
|
if sharpen != 0:
|
|
processed_image = self.apply_image_preprocess(processed_image, sharpen)
|
|
try:
|
|
if crop == "mask_area" and mask is not None:
|
|
processed_image, mask = self.crop_to_mask_area(processed_image, mask)
|
|
clip_vision_output = clip_vision.encode_image(processed_image)
|
|
else:
|
|
crop_bool = crop == "center"
|
|
clip_vision_output = clip_vision.encode_image(processed_image, crop=crop_bool)
|
|
except Exception as e:
|
|
print(f"CLIP视觉编码出错: {e}")
|
|
clip_vision_output = clip_vision.encode_image(processed_image)
|
|
|
|
cond = style_model.get_cond(clip_vision_output)
|
|
|
|
B = cond.shape[0]
|
|
H = W = int(math.sqrt(cond.shape[1]))
|
|
C = cond.shape[2]
|
|
cond = cond.reshape(B, H, W, C)
|
|
|
|
new_H = H * patch_res // 16
|
|
new_W = W * patch_res // 16
|
|
|
|
cond = torch.nn.functional.interpolate(
|
|
cond.permute(0, 3, 1, 2),
|
|
size=(new_H, new_W),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
cond = cond.permute(0, 2, 3, 1)
|
|
cond = cond.reshape(B, -1, C)
|
|
cond = cond.flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
|
|
|
|
c_out = []
|
|
for t in conditioning:
|
|
txt, keys = t
|
|
keys = keys.copy()
|
|
|
|
if prompt_strength != 1.0:
|
|
txt_enhanced = txt * (prompt_strength ** 3)
|
|
txt_repeated = txt_enhanced.repeat(1, 2, 1)
|
|
txt = txt_repeated
|
|
|
|
if style_strength != 1.0:
|
|
processed_cond = self.apply_style_strength(
|
|
cond, txt, style_strength, blend_mode
|
|
)
|
|
else:
|
|
processed_cond = cond
|
|
|
|
if mask is not None:
|
|
feature_size = int(math.sqrt(processed_cond.shape[1]))
|
|
processed_mask = torch.nn.functional.interpolate(
|
|
mask.unsqueeze(1) if mask.dim() == 3 else mask,
|
|
size=(feature_size, feature_size),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).flatten(1).unsqueeze(-1)
|
|
|
|
if txt.shape[1] != processed_cond.shape[1]:
|
|
txt_mean = txt.mean(dim=1, keepdim=True)
|
|
txt_expanded = txt_mean.expand(-1, processed_cond.shape[1], -1)
|
|
else:
|
|
txt_expanded = txt
|
|
|
|
processed_cond = processed_cond * processed_mask + \
|
|
txt_expanded * (1 - processed_mask)
|
|
|
|
if noise_level > 0:
|
|
noise = torch.randn_like(processed_cond)
|
|
noise = (noise - noise.mean()) / (noise.std() + 1e-8)
|
|
processed_cond = torch.lerp(processed_cond, noise, noise_level)
|
|
processed_cond = processed_cond * (1.0 + noise_level)
|
|
|
|
c_out.append([torch.cat((txt, processed_cond), dim=1), keys])
|
|
|
|
return (c_out,)
|
|
|
|
|
|
|