482 lines
19 KiB
Python
482 lines
19 KiB
Python
# Copyright (c) 2024 Sousakujikken HIRO
|
|
# This software is released under the MIT License.
|
|
# https://opensource.org/licenses/MIT
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
class CRTLikeEffectNode:
|
|
kernel_cache = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"gaussian_width_x": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1}),
|
|
"gaussian_width_y": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1}),
|
|
"intensity_scale": ("FLOAT", {"default": 5.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"gamma": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}),
|
|
"gaussian_kernel_size": ([5, 7, 9, 11, 13, 15], {"default": 11}),
|
|
"enable_resize": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}),
|
|
"resize_pixels": ([128, 192, 256, 320, 384, 448, 512], {"default": 256}),
|
|
"max_batch_size": ([1, 2, 4, 8, 16, 24, 32, 48, 64], {"default": 4}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_crt_effect"
|
|
CATEGORY = "image/Pixydust Quantizer🧚✨"
|
|
|
|
def get_gaussian_kernel(self, kernel_size, width_x, width_y, gamma, device, dtype):
|
|
key = (kernel_size, width_x, width_y, gamma)
|
|
if key not in self.kernel_cache:
|
|
x = torch.arange(kernel_size, device=device, dtype=dtype).float()
|
|
y = torch.arange(kernel_size, device=device, dtype=dtype).float()
|
|
Y, X = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
distance_squared = (((X - (kernel_size / 2 - 0.5)) / width_x) ** 2 +
|
|
((Y - (kernel_size / 2 - 0.5)) / width_y) ** 2)
|
|
gaussian = torch.exp(-distance_squared / 2)
|
|
gaussian = torch.pow(gaussian, 1 / gamma)
|
|
gaussian = gaussian / gaussian.sum()
|
|
gaussian = gaussian.unsqueeze(0).unsqueeze(0)
|
|
|
|
self.kernel_cache[key] = gaussian.to(device=device, dtype=dtype)
|
|
return self.kernel_cache[key]
|
|
|
|
def process_batch(self, batch, gaussian_kernel, device, intensity_scale, gamma):
|
|
"""
|
|
単一のミニバッチを処理。データ型の一貫性を保証。
|
|
"""
|
|
# バッチとカーネルを同じデータ型(half)に変換
|
|
batch = batch.half()
|
|
gaussian_kernel = gaussian_kernel.half() # カーネルもhalf精度に変換
|
|
|
|
B, C, H, W = batch.shape
|
|
|
|
# 4倍にアップスケール
|
|
upscaled = F.interpolate(batch, scale_factor=4, mode='nearest')
|
|
_, _, H4, W4 = upscaled.shape
|
|
|
|
# インデックス生成（bool型のままで問題ない）
|
|
rows = torch.arange(H4, device=device).view(1, 1, H4, 1) % 4
|
|
cols = torch.arange(W4, device=device).view(1, 1, 1, W4) % 4
|
|
rows = rows.expand(B, C, H4, W4)
|
|
cols = cols.expand(B, C, H4, W4)
|
|
|
|
# マスク生成（half精度に変換）
|
|
R_mask = ((cols == 0) & (rows < 3)).half()
|
|
G_mask = ((cols == 1) & (rows < 3)).half()
|
|
B_mask = ((cols == 2) & (rows < 3)).half()
|
|
|
|
# シフト処理
|
|
shifted = torch.zeros_like(upscaled)
|
|
shifted[:, 0, :, :] = upscaled[:, 0, :, :] * R_mask[:, 0, :, :]
|
|
shifted[:, 1, :, :] = upscaled[:, 1, :, :] * G_mask[:, 1, :, :]
|
|
shifted[:, 2, :, :] = upscaled[:, 2, :, :] * B_mask[:, 2, :, :]
|
|
|
|
# 畳み込み適用
|
|
convolved = F.conv2d(shifted, gaussian_kernel, padding=0, groups=C)
|
|
|
|
# 後処理
|
|
convolved = convolved * intensity_scale
|
|
convolved = torch.pow(convolved, 1 / gamma)
|
|
convolved = torch.clamp(convolved, 0.0, 1.0)
|
|
|
|
# メモリ解放
|
|
del shifted, upscaled, R_mask, G_mask, B_mask
|
|
torch.cuda.empty_cache()
|
|
|
|
return convolved.permute(0, 2, 3, 1).float()
|
|
|
|
def apply_crt_effect_gpu_optimized(self, image, gaussian_width_x, gaussian_width_y, intensity_scale, gaussian_kernel_size, gamma, max_batch_size):
|
|
"""
|
|
バッチサイズを制限してGPUメモリを効率的に使用
|
|
データ型の一貫性を保証
|
|
"""
|
|
device = image.device
|
|
B, C, H, W = image.shape
|
|
print(f"\nStarting batch processing: total batch size = {B}")
|
|
|
|
# ガウシアンカーネルを準備（まだfloat32のまま）
|
|
gaussian = self.get_gaussian_kernel(gaussian_kernel_size, gaussian_width_x, gaussian_width_y, gamma, device, torch.float32)
|
|
gaussian_kernel = gaussian.repeat(C, 1, 1, 1)
|
|
|
|
# バッチを分割して処理
|
|
outputs = []
|
|
for i in range(0, B, max_batch_size):
|
|
batch = image[i:i + max_batch_size]
|
|
current_batch_size = batch.shape[0]
|
|
print(f"\nProcessing mini-batch {i//max_batch_size + 1}: size = {current_batch_size}")
|
|
|
|
try:
|
|
processed_batch = self.process_batch(
|
|
batch, gaussian_kernel, device,
|
|
intensity_scale, gamma
|
|
)
|
|
print(f"Mini-batch {i//max_batch_size + 1} processed: shape = {processed_batch.shape}")
|
|
|
|
outputs.append(processed_batch.cpu())
|
|
print(f"Mini-batch {i//max_batch_size + 1} transferred to CPU")
|
|
|
|
del processed_batch
|
|
torch.cuda.empty_cache()
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e):
|
|
print(f"OOM detected, processing images individually in mini-batch {i//max_batch_size + 1}")
|
|
for j, single_image in enumerate(batch):
|
|
processed_single = self.process_batch(
|
|
single_image.unsqueeze(0),
|
|
gaussian_kernel, device,
|
|
intensity_scale, gamma
|
|
)
|
|
print(f"Processed individual image {j+1}/{current_batch_size} in mini-batch {i//max_batch_size + 1}")
|
|
outputs.append(processed_single.cpu())
|
|
del processed_single
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
raise e
|
|
|
|
# 全バッチの結果を結合
|
|
final_output = torch.cat(outputs, dim=0)
|
|
print(f"\nAll mini-batches combined: final shape = {final_output.shape}")
|
|
return (final_output,)
|
|
|
|
def apply_crt_effect(self, image, gaussian_width_x, gaussian_width_y, intensity_scale, gamma,
|
|
gaussian_kernel_size, enable_resize, resize_pixels, max_batch_size):
|
|
"""
|
|
メインの処理関数
|
|
"""
|
|
print(f"Input image shape: {image.shape}, dtype: {image.dtype}")
|
|
|
|
if image.max() > 1.0:
|
|
image = image / 255.0
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
image = image.to(device, dtype=torch.float32)
|
|
|
|
B, H, W, C = image.shape
|
|
if C != 3:
|
|
raise ValueError(f"Expected image with 3 channels (RGB), but got {C} channels.")
|
|
|
|
# [B, C, H, W]形式に変換
|
|
image = image.permute(0, 3, 1, 2)
|
|
|
|
# リサイズ処理
|
|
if enable_resize:
|
|
if H > W:
|
|
new_h, new_w = resize_pixels, int(resize_pixels * W / H)
|
|
else:
|
|
new_h, new_w = int(resize_pixels * H / W), resize_pixels
|
|
|
|
try:
|
|
image = F.interpolate(image, size=(new_h, new_w), mode='nearest')
|
|
print(f"After resize: {image.shape}")
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e):
|
|
# リサイズでメモリ不足の場合、バッチごとに処理
|
|
resized_batches = []
|
|
for i in range(0, B, max_batch_size):
|
|
batch = image[i:i + max_batch_size]
|
|
resized_batch = F.interpolate(batch, size=(new_h, new_w), mode='nearest')
|
|
resized_batches.append(resized_batch)
|
|
del batch
|
|
torch.cuda.empty_cache()
|
|
image = torch.cat(resized_batches, dim=0)
|
|
else:
|
|
raise e
|
|
else:
|
|
print(f"No resizing applied. Shape: {image.shape}")
|
|
|
|
# CRT効果の適用
|
|
return self.apply_crt_effect_gpu_optimized(
|
|
image, gaussian_width_x, gaussian_width_y,
|
|
intensity_scale, gaussian_kernel_size, gamma,
|
|
max_batch_size
|
|
)
|
|
|
|
class XYBlurNode:
|
|
kernel_cache = {} # キャッシュを追加してカーネル再計算を防ぐ
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"sigma_x": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"sigma_y": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"max_batch_size": ([1, 2, 4, 8, 16, 24, 32], {"default": 4}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_blur"
|
|
CATEGORY = "image/Pixydust Quantizer🧚✨"
|
|
|
|
def create_gaussian_kernel(self, kernel_size, sigma_x, sigma_y, device, dtype):
|
|
"""
|
|
Generates an anisotropic Gaussian kernel with independent vertical and horizontal components.
|
|
Implements caching for performance optimization.
|
|
"""
|
|
cache_key = (kernel_size, sigma_x, sigma_y, dtype)
|
|
if cache_key not in self.kernel_cache:
|
|
# Generate grid coordinates
|
|
x = torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, dtype=torch.float32, device=device)
|
|
y = torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, dtype=torch.float32, device=device)
|
|
y, x = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
# Calculate Gaussian function
|
|
gaussian = torch.exp(-(x**2 / (2 * sigma_x**2) + y**2 / (2 * sigma_y**2)))
|
|
|
|
# Normalize
|
|
gaussian = gaussian / gaussian.sum()
|
|
|
|
# Cache the kernel
|
|
self.kernel_cache[cache_key] = gaussian.unsqueeze(0).unsqueeze(0).to(dtype=dtype)
|
|
|
|
return self.kernel_cache[cache_key]
|
|
|
|
def calculate_kernel_size(self, sigma_x, sigma_y, max_kernel_size=31):
|
|
"""
|
|
Calculates appropriate kernel size based on sigma values.
|
|
"""
|
|
max_sigma = max(sigma_x, sigma_y)
|
|
kernel_size = int(6 * max_sigma + 1)
|
|
kernel_size = max(3, min(kernel_size + (kernel_size % 2 == 0), max_kernel_size))
|
|
return kernel_size
|
|
|
|
def process_batch(self, batch, kernel, pad_size, intensity, device):
|
|
"""
|
|
Process a single batch of images with the blur effect.
|
|
"""
|
|
# Ensure correct data type (half precision for GPU efficiency)
|
|
batch = batch.half()
|
|
kernel = kernel.half()
|
|
|
|
# Apply convolution
|
|
blurred = F.conv2d(
|
|
batch,
|
|
kernel,
|
|
padding=pad_size,
|
|
groups=batch.shape[1] # Process each channel independently
|
|
)
|
|
|
|
# Apply intensity blending
|
|
blurred = batch * (1 - intensity) + blurred * intensity
|
|
|
|
# Clamp values
|
|
blurred = torch.clamp(blurred, 0, 1)
|
|
|
|
# Convert back to float32 for output
|
|
return blurred.float()
|
|
|
|
def apply_blur_optimized(self, image, sigma_x, sigma_y, intensity, max_batch_size):
|
|
"""
|
|
Applies blur effect with optimized batch processing and GPU memory management.
|
|
"""
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
B, C, H, W = image.shape
|
|
print(f"\nProcessing total batch size: {B}")
|
|
|
|
# Calculate kernel size and create kernel
|
|
kernel_size = self.calculate_kernel_size(sigma_x, sigma_y)
|
|
kernel = self.create_gaussian_kernel(kernel_size, sigma_x, sigma_y, device, torch.float32)
|
|
kernel = kernel.repeat(C, 1, 1, 1)
|
|
pad_size = kernel_size // 2
|
|
|
|
# Process batches
|
|
outputs = []
|
|
for i in range(0, B, max_batch_size):
|
|
batch = image[i:i + max_batch_size]
|
|
current_batch_size = batch.shape[0]
|
|
print(f"\nProcessing mini-batch {i//max_batch_size + 1}: size = {current_batch_size}")
|
|
|
|
try:
|
|
processed_batch = self.process_batch(
|
|
batch, kernel, pad_size, intensity, device
|
|
)
|
|
print(f"Mini-batch {i//max_batch_size + 1} processed: shape = {processed_batch.shape}")
|
|
|
|
# Move processed batch to CPU to free GPU memory
|
|
outputs.append(processed_batch.cpu())
|
|
print(f"Mini-batch {i//max_batch_size + 1} transferred to CPU")
|
|
|
|
del processed_batch
|
|
torch.cuda.empty_cache()
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e):
|
|
print(f"OOM detected, processing images individually in mini-batch {i//max_batch_size + 1}")
|
|
# Process images one by one if batch processing fails
|
|
for j, single_image in enumerate(batch):
|
|
processed_single = self.process_batch(
|
|
single_image.unsqueeze(0),
|
|
kernel, pad_size, intensity, device
|
|
)
|
|
print(f"Processed individual image {j+1}/{current_batch_size} in mini-batch {i//max_batch_size + 1}")
|
|
outputs.append(processed_single.cpu())
|
|
del processed_single
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
raise e
|
|
|
|
# Combine all processed batches
|
|
final_output = torch.cat(outputs, dim=0)
|
|
print(f"\nAll mini-batches combined: final shape = {final_output.shape}")
|
|
return final_output
|
|
|
|
def apply_blur(self, image: torch.Tensor, sigma_x: float, sigma_y: float,
|
|
intensity: float, max_batch_size: int):
|
|
"""
|
|
Main entry point for the blur effect.
|
|
"""
|
|
print(f"Input image shape: {image.shape}, dtype: {image.dtype}")
|
|
|
|
# Normalize input if necessary
|
|
if image.max() > 1.0:
|
|
image = image / 255.0
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
image = image.to(device)
|
|
|
|
# Convert from [B, H, W, C] to [B, C, H, W]
|
|
image = image.permute(0, 3, 1, 2)
|
|
|
|
# Apply optimized blur processing
|
|
blurred = self.apply_blur_optimized(
|
|
image, sigma_x, sigma_y, intensity, max_batch_size
|
|
)
|
|
|
|
# Convert back to [B, H, W, C] format
|
|
blurred = blurred.permute(0, 2, 3, 1)
|
|
|
|
# Print debug information
|
|
print(f"After blur: mean={blurred.mean().item():.4f}, "
|
|
f"std={blurred.std().item():.4f}, "
|
|
f"min={blurred.min().item():.4f}, "
|
|
f"max={blurred.max().item():.4f}")
|
|
|
|
return (blurred,)
|
|
|
|
# class XYBlurNode:
|
|
# @classmethod
|
|
# def INPUT_TYPES(cls):
|
|
# return {
|
|
# "required": {
|
|
# "image": ("IMAGE",),
|
|
# "sigma_x": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
# "sigma_y": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
# "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
# }
|
|
# }
|
|
|
|
# RETURN_TYPES = ("IMAGE",)
|
|
# FUNCTION = "apply_blur"
|
|
# CATEGORY = "image/Pixydust Quantizer🧚✨"
|
|
|
|
# def create_gaussian_kernel(self, kernel_size, sigma_x, sigma_y, device, dtype):
|
|
# """
|
|
# Generates an anisotropic Gaussian kernel with independent vertical and horizontal components.
|
|
# """
|
|
# # Generate grid coordinates
|
|
# x = torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, dtype=torch.float32, device=device)
|
|
# y = torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, dtype=torch.float32, device=device)
|
|
# y, x = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
# # Calculate Gaussian function (independent for vertical and horizontal)
|
|
# gaussian = torch.exp(-(x**2 / (2 * sigma_x**2) + y**2 / (2 * sigma_y**2)))
|
|
|
|
# # Normalize so that the sum is 1
|
|
# gaussian = gaussian / gaussian.sum()
|
|
|
|
# # Adjust the shape of the kernel to [1, 1, kernel_size, kernel_size]
|
|
# return gaussian.unsqueeze(0).unsqueeze(0).to(dtype=dtype)
|
|
|
|
# def calculate_kernel_size(self, sigma_x, sigma_y, max_kernel_size=31):
|
|
# """
|
|
# Calculates an appropriate kernel size based on sigma_x and sigma_y.
|
|
# Typically, the kernel size is set to 6 * sigma and rounded to an odd number.
|
|
# """
|
|
# # Calculate kernel size based on the largest sigma
|
|
# max_sigma = max(sigma_x, sigma_y)
|
|
# kernel_size = int(6 * max_sigma + 1)
|
|
|
|
# # Adjust kernel size to be odd
|
|
# if kernel_size % 2 == 0:
|
|
# kernel_size += 1
|
|
|
|
# # Set minimum kernel size to 3
|
|
# kernel_size = max(kernel_size, 3)
|
|
|
|
# # Set maximum kernel size
|
|
# kernel_size = min(kernel_size, max_kernel_size)
|
|
|
|
# return kernel_size
|
|
|
|
# def apply_blur(self, image: torch.Tensor, sigma_x: float, sigma_y: float, intensity: float):
|
|
# """
|
|
# Applies an anisotropic Gaussian blur to the image.
|
|
# The kernel size is automatically calculated based on sigma_x and sigma_y.
|
|
# """
|
|
# device = image.device
|
|
# dtype = image.dtype
|
|
# B, H, W, C = image.shape
|
|
|
|
# # Calculate kernel size
|
|
# kernel_size = self.calculate_kernel_size(sigma_x, sigma_y)
|
|
# print(f"Calculated kernel_size: {kernel_size} based on sigma_x: {sigma_x}, sigma_y: {sigma_y}")
|
|
|
|
# # Convert to [B, C, H, W] format
|
|
# x = image.permute(0, 3, 1, 2)
|
|
|
|
# # Generate Gaussian kernel
|
|
# kernel = self.create_gaussian_kernel(kernel_size, sigma_x, sigma_y, device, dtype)
|
|
|
|
# # Repeat the same kernel for each channel
|
|
# kernel = kernel.repeat(C, 1, 1, 1)
|
|
|
|
# # Calculate padding size
|
|
# pad_size = kernel_size // 2
|
|
|
|
# # Apply convolution operation (independently for each channel)
|
|
# blurred = F.conv2d(
|
|
# x,
|
|
# kernel,
|
|
# padding=pad_size,
|
|
# groups=C # Process each channel independently
|
|
# )
|
|
|
|
# # Apply the intensity parameter
|
|
# # Allows intensity to be in the range [0, 2]
|
|
# # intensity = 1.0 for normal blur
|
|
# # intensity < 1.0 to decrease the blur intensity
|
|
# # intensity > 1.0 to increase the blur intensity
|
|
# blurred = x * (1 - intensity) + blurred * intensity
|
|
|
|
# # Convert back to [B, H, W, C] format
|
|
# blurred = blurred.permute(0, 2, 3, 1)
|
|
|
|
# # Clamp the values to [0, 1]
|
|
# blurred = torch.clamp(blurred, 0, 1)
|
|
|
|
# # Output debug information
|
|
# print(f"After blur: mean={blurred.mean().item():.4f}, std={blurred.std().item():.4f}, min={blurred.min().item():.4f}, max={blurred.max().item():.4f}")
|
|
|
|
# # Move the tensor to CPU before returning
|
|
# blurred = blurred.cpu()
|
|
|
|
# return (blurred,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CRTLikeEffectNode": CRTLikeEffectNode,
|
|
"XYBlurNode": XYBlurNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CRTLikeEffectNode": "CRTLike Effect",
|
|
"XYBlurNode": "XY Blur"
|
|
}
|