effects
This commit is contained in:
Dag Thomas Olsen
+82
@@ -217,6 +217,88 @@ try:
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular UniversalVisionCloner: {e}")
|
||||
|
||||
# APNext FX Nodes
|
||||
try:
|
||||
from .nodes.image_fx.bloom import APNextBloom
|
||||
NEW_MAPPINGS["APNextBloom"] = APNextBloom
|
||||
NEW_DISPLAY_MAPPINGS["APNextBloom"] = "APNext Bloom FX"
|
||||
# print("✅ Loaded modular APNextBloom")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextBloom: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.sharpen import APNextSharpen
|
||||
NEW_MAPPINGS["APNextSharpen"] = APNextSharpen
|
||||
NEW_DISPLAY_MAPPINGS["APNextSharpen"] = "APNext Sharpen FX"
|
||||
# print("✅ Loaded modular APNextSharpen")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextSharpen: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.noise import APNextNoise
|
||||
NEW_MAPPINGS["APNextNoise"] = APNextNoise
|
||||
NEW_DISPLAY_MAPPINGS["APNextNoise"] = "APNext Noise FX"
|
||||
# print("✅ Loaded modular APNextNoise")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextNoise: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.rough import APNextRough
|
||||
NEW_MAPPINGS["APNextRough"] = APNextRough
|
||||
NEW_DISPLAY_MAPPINGS["APNextRough"] = "APNext Rough FX"
|
||||
# print("✅ Loaded modular APNextRough")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextRough: {e}")
|
||||
|
||||
# APNext Advanced FX Nodes
|
||||
try:
|
||||
from .nodes.image_fx.color_grading import APNextColorGrading
|
||||
NEW_MAPPINGS["APNextColorGrading"] = APNextColorGrading
|
||||
NEW_DISPLAY_MAPPINGS["APNextColorGrading"] = "APNext Color Grading FX"
|
||||
# print("✅ Loaded modular APNextColorGrading")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextColorGrading: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.cross_processing import APNextCrossProcessing
|
||||
NEW_MAPPINGS["APNextCrossProcessing"] = APNextCrossProcessing
|
||||
NEW_DISPLAY_MAPPINGS["APNextCrossProcessing"] = "APNext Cross Processing FX"
|
||||
# print("✅ Loaded modular APNextCrossProcessing")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextCrossProcessing: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.split_toning import APNextSplitToning
|
||||
NEW_MAPPINGS["APNextSplitToning"] = APNextSplitToning
|
||||
NEW_DISPLAY_MAPPINGS["APNextSplitToning"] = "APNext Split Toning FX"
|
||||
# print("✅ Loaded modular APNextSplitToning")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextSplitToning: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.hdr_tone_mapping import APNextHDRToneMapping
|
||||
NEW_MAPPINGS["APNextHDRToneMapping"] = APNextHDRToneMapping
|
||||
NEW_DISPLAY_MAPPINGS["APNextHDRToneMapping"] = "APNext HDR Tone Mapping FX"
|
||||
# print("✅ Loaded modular APNextHDRToneMapping")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextHDRToneMapping: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.glitch_art import APNextGlitchArt
|
||||
NEW_MAPPINGS["APNextGlitchArt"] = APNextGlitchArt
|
||||
NEW_DISPLAY_MAPPINGS["APNextGlitchArt"] = "APNext Glitch Art FX"
|
||||
# print("✅ Loaded modular APNextGlitchArt")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextGlitchArt: {e}")
|
||||
|
||||
try:
|
||||
from .nodes.image_fx.film_halation import APNextFilmHalation
|
||||
NEW_MAPPINGS["APNextFilmHalation"] = APNextFilmHalation
|
||||
NEW_DISPLAY_MAPPINGS["APNextFilmHalation"] = "APNext Film Halation FX"
|
||||
# print("✅ Loaded modular APNextFilmHalation")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not load modular APNextFilmHalation: {e}")
|
||||
|
||||
# Combine mappings (modular nodes + dynamic nodes)
|
||||
NODE_CLASS_MAPPINGS = {**NEW_MAPPINGS, **DYNAMIC_MAPPINGS}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**NEW_DISPLAY_MAPPINGS, **DYNAMIC_DISPLAY_MAPPINGS}
|
||||
|
||||
@@ -1 +1,11 @@
|
||||
# Node modules for comfyui_dagthomas
|
||||
|
||||
# Import FX nodes
|
||||
from .image_fx import NODE_CLASS_MAPPINGS as FX_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FX_NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
# Add FX nodes to the global mappings
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(FX_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(FX_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
@@ -0,0 +1,38 @@
|
||||
# APNext FX Image Effects Nodes
|
||||
|
||||
from .bloom import APNextBloom
|
||||
from .sharpen import APNextSharpen
|
||||
from .noise import APNextNoise
|
||||
from .rough import APNextRough
|
||||
from .color_grading import APNextColorGrading
|
||||
from .cross_processing import APNextCrossProcessing
|
||||
from .split_toning import APNextSplitToning
|
||||
from .hdr_tone_mapping import APNextHDRToneMapping
|
||||
from .glitch_art import APNextGlitchArt
|
||||
from .film_halation import APNextFilmHalation
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"APNextBloom": APNextBloom,
|
||||
"APNextSharpen": APNextSharpen,
|
||||
"APNextNoise": APNextNoise,
|
||||
"APNextRough": APNextRough,
|
||||
"APNextColorGrading": APNextColorGrading,
|
||||
"APNextCrossProcessing": APNextCrossProcessing,
|
||||
"APNextSplitToning": APNextSplitToning,
|
||||
"APNextHDRToneMapping": APNextHDRToneMapping,
|
||||
"APNextGlitchArt": APNextGlitchArt,
|
||||
"APNextFilmHalation": APNextFilmHalation,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"APNextBloom": "APNext Bloom FX",
|
||||
"APNextSharpen": "APNext Sharpen FX",
|
||||
"APNextNoise": "APNext Noise FX",
|
||||
"APNextRough": "APNext Rough FX",
|
||||
"APNextColorGrading": "APNext Color Grading FX",
|
||||
"APNextCrossProcessing": "APNext Cross Processing FX",
|
||||
"APNextSplitToning": "APNext Split Toning FX",
|
||||
"APNextHDRToneMapping": "APNext HDR Tone Mapping FX",
|
||||
"APNextGlitchArt": "APNext Glitch Art FX",
|
||||
"APNextFilmHalation": "APNext Film Halation FX",
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
# APNext Bloom Effect Node
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextBloom:
|
||||
"""
|
||||
APNext Bloom Effect Node
|
||||
Creates a bloom effect by making bright areas glow using gaussian blur and blend modes
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
||||
"threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"blur_radius": ("FLOAT", {"default": 15.0, "min": 1.0, "max": 50.0, "step": 0.5}),
|
||||
"blend_mode": (["additive", "screen", "overlay"], {"default": "additive"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_bloom"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX"
|
||||
|
||||
def apply_bloom(self, images, intensity, threshold, blur_radius, blend_mode):
|
||||
"""Apply bloom effect to images using optimized tensor operations"""
|
||||
# Work directly with tensors for better performance
|
||||
batch_size = images.shape[0]
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Process all images in batch
|
||||
result = self._create_bloom_effect_tensor(
|
||||
images, intensity, threshold, blur_radius, blend_mode, device, dtype
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _create_bloom_effect_tensor(self, images, intensity, threshold, blur_radius, blend_mode, device, dtype):
|
||||
"""Create bloom effect using tensor operations for better performance"""
|
||||
# Apply threshold to isolate bright areas
|
||||
bloom_layer = self._apply_threshold_tensor(images, threshold)
|
||||
|
||||
# Apply gaussian blur using torch operations
|
||||
bloom_layer = self._gaussian_blur_tensor(bloom_layer, blur_radius)
|
||||
|
||||
# Blend the bloom layer with the original image
|
||||
result = self._blend_images_tensor(images, bloom_layer, blend_mode, intensity)
|
||||
|
||||
return result
|
||||
|
||||
def _apply_threshold_tensor(self, images, threshold):
|
||||
"""Apply threshold to isolate bright areas using tensor operations"""
|
||||
# Calculate luminance using tensor operations
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
|
||||
# Create mask for bright areas
|
||||
bright_mask = luminance > threshold
|
||||
|
||||
# Apply threshold - keep only pixels above threshold
|
||||
result = images.clone()
|
||||
bright_mask = bright_mask.unsqueeze(-1).expand_as(images)
|
||||
result = torch.where(bright_mask, images, torch.zeros_like(images))
|
||||
|
||||
return result
|
||||
|
||||
def _blend_images_tensor(self, base, overlay, blend_mode, intensity):
|
||||
"""Blend two images using tensor operations"""
|
||||
overlay = overlay * intensity
|
||||
|
||||
if blend_mode == "additive":
|
||||
result = base + overlay
|
||||
elif blend_mode == "screen":
|
||||
result = 1 - (1 - base) * (1 - overlay)
|
||||
elif blend_mode == "overlay":
|
||||
mask = base < 0.5
|
||||
result = torch.where(
|
||||
mask,
|
||||
2 * base * overlay,
|
||||
1 - 2 * (1 - base) * (1 - overlay)
|
||||
)
|
||||
|
||||
# Clamp values to [0, 1]
|
||||
result = torch.clamp(result, 0, 1)
|
||||
return result
|
||||
|
||||
def _gaussian_blur_tensor(self, images, radius):
|
||||
"""Apply gaussian blur using torch operations"""
|
||||
# Convert radius to sigma (approximation)
|
||||
sigma = radius / 3.0
|
||||
|
||||
# Create gaussian kernel
|
||||
kernel_size = int(2 * radius + 1)
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
|
||||
# Create 1D gaussian kernel
|
||||
x = torch.arange(kernel_size, dtype=images.dtype, device=images.device)
|
||||
x = x - kernel_size // 2
|
||||
kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2)
|
||||
kernel_1d = kernel_1d / kernel_1d.sum()
|
||||
|
||||
# Apply separable gaussian blur
|
||||
# Reshape for conv operations: [batch, channels, height, width]
|
||||
images_conv = images.permute(0, 3, 1, 2)
|
||||
|
||||
# Horizontal blur
|
||||
kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1)
|
||||
blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3)
|
||||
|
||||
# Vertical blur
|
||||
kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1)
|
||||
blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3)
|
||||
|
||||
# Reshape back to original format: [batch, height, width, channels]
|
||||
return blurred.permute(0, 2, 3, 1)
|
||||
@@ -0,0 +1,533 @@
|
||||
# APNext Color Grading Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import os
|
||||
import glob
|
||||
import folder_paths
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
from ...utils.image_utils import tensor2pil, pil2tensor
|
||||
|
||||
|
||||
class APNextColorGrading:
|
||||
"""
|
||||
APNext Color Grading Effect Node
|
||||
Applies color grading using LUT files or manual controls
|
||||
|
||||
Supported LUT formats:
|
||||
- .cube files (Adobe/Blackmagic)
|
||||
- .3dl files (Autodesk/Flame)
|
||||
- Image LUTs (.png, .jpg, .tiff, .exr)
|
||||
|
||||
Example paths:
|
||||
- C:/LUTs/cinematic.cube
|
||||
- ./assets/vintage_film.3dl
|
||||
- /path/to/lut_image.png
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# Register LUT file types with ComfyUI's folder_paths system
|
||||
cls._register_lut_paths()
|
||||
|
||||
# Get LUT files using ComfyUI's standard system
|
||||
lut_files = folder_paths.get_filename_list("luts")
|
||||
if not lut_files:
|
||||
lut_files = ["None"]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"method": (["manual", "lut_file"], {"default": "manual"}),
|
||||
"lut_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
# Manual controls
|
||||
"exposure": ("FLOAT", {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.01}),
|
||||
"contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}),
|
||||
"highlights": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
"shadows": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
"saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}),
|
||||
"temperature": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
"tint": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
"optional": {
|
||||
"lut_file": (lut_files, {"default": "None"}),
|
||||
"custom_lut_path": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Custom LUT path (overrides picker)"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _register_lut_paths(cls):
|
||||
"""Register LUT file paths with ComfyUI's folder system"""
|
||||
try:
|
||||
# Add LUT directory to ComfyUI's folder paths
|
||||
if "luts" not in folder_paths.folder_names_and_paths:
|
||||
lut_dirs = [
|
||||
os.path.join(folder_paths.models_dir, "luts"),
|
||||
os.path.join(folder_paths.input_directory, "luts"),
|
||||
"./luts",
|
||||
"./LUTs"
|
||||
]
|
||||
|
||||
# Create luts directory in models if it doesn't exist
|
||||
main_lut_dir = os.path.join(folder_paths.models_dir, "luts")
|
||||
if not os.path.exists(main_lut_dir):
|
||||
os.makedirs(main_lut_dir, exist_ok=True)
|
||||
print(f"📁 Created LUT directory: {main_lut_dir}")
|
||||
|
||||
# Register with ComfyUI
|
||||
folder_paths.folder_names_and_paths["luts"] = (lut_dirs, {".cube", ".3dl", ".png", ".jpg", ".jpeg", ".tiff", ".exr"})
|
||||
print("✅ Registered LUT file paths with ComfyUI")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not register LUT paths: {e}")
|
||||
# Fallback to manual scanning
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def _scan_for_lut_files(cls):
|
||||
"""Scan common directories for LUT files"""
|
||||
lut_files = ["None"] # Default option
|
||||
|
||||
# Common LUT directories to scan
|
||||
search_dirs = [
|
||||
"./luts",
|
||||
"./LUTs",
|
||||
"./assets/luts",
|
||||
"./custom_nodes/comfyui_dagthomas/luts",
|
||||
os.path.expanduser("~/LUTs"),
|
||||
os.path.expanduser("~/Documents/LUTs"),
|
||||
"C:/LUTs" if os.name == 'nt' else "/usr/local/share/luts",
|
||||
]
|
||||
|
||||
# LUT file extensions
|
||||
extensions = ['*.cube', '*.3dl', '*.png', '*.jpg', '*.jpeg', '*.tiff', '*.exr']
|
||||
|
||||
for directory in search_dirs:
|
||||
if os.path.exists(directory):
|
||||
for ext in extensions:
|
||||
pattern = os.path.join(directory, '**', ext)
|
||||
files = glob.glob(pattern, recursive=True)
|
||||
for file_path in files:
|
||||
# Create a display name (relative path or just filename)
|
||||
if len(file_path) > 60:
|
||||
display_name = f"...{file_path[-57:]}"
|
||||
else:
|
||||
display_name = file_path
|
||||
lut_files.append(f"{display_name}|{file_path}")
|
||||
|
||||
# Also scan current working directory
|
||||
for ext in extensions:
|
||||
files = glob.glob(ext)
|
||||
for file_path in files:
|
||||
lut_files.append(f"{file_path}|{os.path.abspath(file_path)}")
|
||||
|
||||
return lut_files
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_color_grading"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_color_grading(self, images, method, lut_strength, exposure, contrast,
|
||||
highlights, shadows, saturation, temperature, tint,
|
||||
lut_file="None", custom_lut_path=""):
|
||||
"""Apply color grading to images"""
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
if method == "lut_file":
|
||||
# Determine which LUT path to use
|
||||
lut_path = None
|
||||
|
||||
# Custom path takes priority
|
||||
if custom_lut_path and custom_lut_path.strip():
|
||||
lut_path = custom_lut_path.strip()
|
||||
print(f"📁 Using custom LUT path: {lut_path}")
|
||||
|
||||
# Otherwise use ComfyUI file selection
|
||||
elif lut_file and lut_file != "None":
|
||||
lut_path = folder_paths.get_full_path("luts", lut_file)
|
||||
print(f"📁 Using selected LUT: {lut_path}")
|
||||
|
||||
# Apply LUT if we have a valid path
|
||||
if lut_path and os.path.exists(lut_path):
|
||||
result = self._apply_lut_grading(images, lut_path, lut_strength)
|
||||
elif lut_path:
|
||||
print(f"⚠️ LUT file not found: {lut_path}")
|
||||
print("Using manual grading instead.")
|
||||
result = self._apply_manual_grading(
|
||||
images, exposure, contrast, highlights, shadows,
|
||||
saturation, temperature, tint
|
||||
)
|
||||
else:
|
||||
print("⚠️ No LUT file selected. Using manual grading instead.")
|
||||
result = self._apply_manual_grading(
|
||||
images, exposure, contrast, highlights, shadows,
|
||||
saturation, temperature, tint
|
||||
)
|
||||
else:
|
||||
result = self._apply_manual_grading(
|
||||
images, exposure, contrast, highlights, shadows,
|
||||
saturation, temperature, tint
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _apply_lut_grading(self, images, lut_path, strength):
|
||||
"""Apply LUT-based color grading"""
|
||||
try:
|
||||
# Load LUT file
|
||||
lut = self._load_lut_file(lut_path)
|
||||
if lut is None:
|
||||
return images
|
||||
|
||||
# Apply LUT to images
|
||||
result = self._apply_lut_tensor(images, lut, strength)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error applying LUT: {e}")
|
||||
return images
|
||||
|
||||
def _apply_manual_grading(self, images, exposure, contrast, highlights, shadows,
|
||||
saturation, temperature, tint):
|
||||
"""Apply manual color grading controls"""
|
||||
result = images.clone()
|
||||
|
||||
# Apply exposure
|
||||
if exposure != 0.0:
|
||||
exposure_factor = 2.0 ** exposure
|
||||
result = result * exposure_factor
|
||||
|
||||
# Apply contrast
|
||||
if contrast != 1.0:
|
||||
# Contrast around middle gray (0.18 in linear, ~0.5 in sRGB)
|
||||
mid_gray = 0.5
|
||||
result = (result - mid_gray) * contrast + mid_gray
|
||||
|
||||
# Apply highlights and shadows
|
||||
if highlights != 0.0 or shadows != 0.0:
|
||||
result = self._apply_highlight_shadow_tensor(result, highlights, shadows)
|
||||
|
||||
# Apply saturation
|
||||
if saturation != 1.0:
|
||||
result = self._apply_saturation_tensor(result, saturation)
|
||||
|
||||
# Apply temperature and tint
|
||||
if temperature != 0.0 or tint != 0.0:
|
||||
result = self._apply_white_balance_tensor(result, temperature, tint)
|
||||
|
||||
# Clamp values
|
||||
result = torch.clamp(result, 0, 1)
|
||||
return result
|
||||
|
||||
def _apply_highlight_shadow_tensor(self, images, highlights, shadows):
|
||||
"""Apply highlight and shadow adjustments using tensor operations"""
|
||||
# Calculate luminance
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Create masks for highlights and shadows
|
||||
highlight_mask = torch.sigmoid((luminance - 0.7) * 10) # Soft transition around 0.7
|
||||
shadow_mask = torch.sigmoid((0.3 - luminance) * 10) # Soft transition around 0.3
|
||||
|
||||
# Apply adjustments
|
||||
highlight_adj = 1.0 + highlights * highlight_mask
|
||||
shadow_adj = 1.0 + shadows * shadow_mask
|
||||
|
||||
result = images * highlight_adj * shadow_adj
|
||||
return result
|
||||
|
||||
def _apply_saturation_tensor(self, images, saturation):
|
||||
"""Apply saturation adjustment using tensor operations"""
|
||||
# Calculate luminance
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Blend between grayscale and original based on saturation
|
||||
result = luminance + (images - luminance) * saturation
|
||||
return result
|
||||
|
||||
def _apply_white_balance_tensor(self, images, temperature, tint):
|
||||
"""Apply white balance (temperature and tint) using tensor operations"""
|
||||
# Temperature affects blue-yellow balance
|
||||
if temperature != 0.0:
|
||||
# Warm up (positive) or cool down (negative)
|
||||
temp_matrix = torch.tensor([
|
||||
[1.0 + temperature * 0.2, 0.0, -temperature * 0.1],
|
||||
[0.0, 1.0, 0.0],
|
||||
[-temperature * 0.1, 0.0, 1.0 - temperature * 0.2]
|
||||
], device=images.device, dtype=images.dtype)
|
||||
|
||||
# Apply color matrix
|
||||
images_flat = images.view(-1, 3)
|
||||
images_adjusted = torch.matmul(images_flat, temp_matrix.t())
|
||||
images = images_adjusted.view(images.shape)
|
||||
|
||||
# Tint affects green-magenta balance
|
||||
if tint != 0.0:
|
||||
tint_matrix = torch.tensor([
|
||||
[1.0 - tint * 0.1, tint * 0.2, 0.0],
|
||||
[-tint * 0.1, 1.0 + tint * 0.1, 0.0],
|
||||
[0.0, 0.0, 1.0]
|
||||
], device=images.device, dtype=images.dtype)
|
||||
|
||||
# Apply color matrix
|
||||
images_flat = images.view(-1, 3)
|
||||
images_adjusted = torch.matmul(images_flat, tint_matrix.t())
|
||||
images = images_adjusted.view(images.shape)
|
||||
|
||||
return images
|
||||
|
||||
def _load_lut_file(self, lut_path):
|
||||
"""Load LUT file (supports .cube, .3dl, and image formats)"""
|
||||
try:
|
||||
file_ext = lut_path.lower()
|
||||
if file_ext.endswith('.cube'):
|
||||
return self._load_cube_lut(lut_path)
|
||||
elif file_ext.endswith('.3dl'):
|
||||
return self._load_3dl_lut(lut_path)
|
||||
elif file_ext.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.exr')):
|
||||
return self._load_image_lut(lut_path)
|
||||
else:
|
||||
print(f"Unsupported LUT format: {lut_path}")
|
||||
print("Supported formats: .cube, .3dl, .png, .jpg, .jpeg, .tiff, .exr")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error loading LUT file {lut_path}: {e}")
|
||||
return None
|
||||
|
||||
def _load_cube_lut(self, cube_path):
|
||||
"""Load .cube format LUT file with enhanced compatibility"""
|
||||
try:
|
||||
# Try different encodings for better compatibility
|
||||
encodings = ['utf-8', 'latin-1', 'cp1252']
|
||||
lines = None
|
||||
|
||||
for encoding in encodings:
|
||||
try:
|
||||
with open(cube_path, 'r', encoding=encoding) as f:
|
||||
lines = f.readlines()
|
||||
break
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
|
||||
if lines is None:
|
||||
print(f"Could not read .cube file with any supported encoding: {cube_path}")
|
||||
return None
|
||||
|
||||
lut_size = 33 # Default size
|
||||
lut_data = []
|
||||
domain_min = [0.0, 0.0, 0.0]
|
||||
domain_max = [1.0, 1.0, 1.0]
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
|
||||
# Skip empty lines and comments
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
# Parse header information
|
||||
if line.startswith('LUT_3D_SIZE'):
|
||||
lut_size = int(line.split()[-1])
|
||||
print(f"Loading .cube LUT with size: {lut_size}x{lut_size}x{lut_size}")
|
||||
|
||||
elif line.startswith('DOMAIN_MIN'):
|
||||
domain_min = [float(x) for x in line.split()[1:4]]
|
||||
|
||||
elif line.startswith('DOMAIN_MAX'):
|
||||
domain_max = [float(x) for x in line.split()[1:4]]
|
||||
|
||||
elif line.startswith('TITLE'):
|
||||
title = line.split('"')[1] if '"' in line else line.split()[1]
|
||||
print(f"Loading LUT: {title}")
|
||||
|
||||
# Parse LUT data
|
||||
elif not any(line.startswith(keyword) for keyword in ['LUT_1D_SIZE', 'LUT_1D_INPUT_RANGE']):
|
||||
try:
|
||||
values = [float(x) for x in line.split()]
|
||||
if len(values) == 3:
|
||||
# Normalize values if they're outside 0-1 range
|
||||
normalized_values = []
|
||||
for i, val in enumerate(values):
|
||||
if domain_max[i] != domain_min[i]:
|
||||
normalized = (val - domain_min[i]) / (domain_max[i] - domain_min[i])
|
||||
else:
|
||||
normalized = val
|
||||
normalized_values.append(max(0.0, min(1.0, normalized)))
|
||||
|
||||
lut_data.append(normalized_values)
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
|
||||
# Validate LUT data
|
||||
expected_size = lut_size ** 3
|
||||
if len(lut_data) == expected_size:
|
||||
lut_array = np.array(lut_data, dtype=np.float32).reshape(lut_size, lut_size, lut_size, 3)
|
||||
print(f"Successfully loaded .cube LUT: {len(lut_data)} entries")
|
||||
return torch.from_numpy(lut_array)
|
||||
else:
|
||||
print(f"Invalid LUT data size: expected {expected_size}, got {len(lut_data)}")
|
||||
|
||||
# Try to handle common size mismatches
|
||||
if len(lut_data) > 0:
|
||||
# Find the closest cube root
|
||||
cube_root = round(len(lut_data) ** (1/3))
|
||||
if cube_root ** 3 == len(lut_data):
|
||||
print(f"Adjusting LUT size to {cube_root}x{cube_root}x{cube_root}")
|
||||
lut_array = np.array(lut_data, dtype=np.float32).reshape(cube_root, cube_root, cube_root, 3)
|
||||
return torch.from_numpy(lut_array)
|
||||
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading .cube file {cube_path}: {e}")
|
||||
return None
|
||||
|
||||
def _load_3dl_lut(self, path_3dl):
|
||||
"""Load .3dl format LUT file"""
|
||||
try:
|
||||
with open(path_3dl, 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
lut_data = []
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#'):
|
||||
try:
|
||||
values = [float(x) for x in line.split()]
|
||||
if len(values) == 3:
|
||||
lut_data.append(values)
|
||||
except:
|
||||
continue
|
||||
|
||||
# 3dl files are typically 32x32x32
|
||||
if len(lut_data) == 32768: # 32^3
|
||||
lut_array = np.array(lut_data, dtype=np.float32).reshape(32, 32, 32, 3)
|
||||
print(f"Successfully loaded .3dl LUT: 32x32x32")
|
||||
return torch.from_numpy(lut_array)
|
||||
else:
|
||||
# Try to determine size
|
||||
cube_root = round(len(lut_data) ** (1/3))
|
||||
if cube_root ** 3 == len(lut_data):
|
||||
lut_array = np.array(lut_data, dtype=np.float32).reshape(cube_root, cube_root, cube_root, 3)
|
||||
print(f"Successfully loaded .3dl LUT: {cube_root}x{cube_root}x{cube_root}")
|
||||
return torch.from_numpy(lut_array)
|
||||
else:
|
||||
print(f"Invalid .3dl LUT size: {len(lut_data)} entries")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading .3dl file {path_3dl}: {e}")
|
||||
return None
|
||||
|
||||
def _load_image_lut(self, image_path):
|
||||
"""Load LUT from image (supports various image formats including Photoshop LUT images)"""
|
||||
try:
|
||||
lut_image = Image.open(image_path).convert('RGB')
|
||||
lut_array = np.array(lut_image, dtype=np.float32) / 255.0
|
||||
height, width = lut_array.shape[:2]
|
||||
|
||||
print(f"Loading image LUT: {width}x{height}")
|
||||
|
||||
# Common LUT image formats:
|
||||
# 512x512 for 64^3 LUT (8x8 grid)
|
||||
# 1024x32 for 32^3 LUT (32x1 strip)
|
||||
# 256x16 for 16^3 LUT (16x1 strip)
|
||||
|
||||
# Try different standard formats
|
||||
formats = [
|
||||
(512, 512, 64), # 64^3 in 8x8 grid
|
||||
(1024, 32, 32), # 32^3 in horizontal strip
|
||||
(256, 16, 16), # 16^3 in horizontal strip
|
||||
(2048, 64, 64), # 64^3 in horizontal strip
|
||||
]
|
||||
|
||||
for expected_w, expected_h, lut_size in formats:
|
||||
if width == expected_w and height == expected_h:
|
||||
if expected_w == expected_h:
|
||||
# Square format (grid layout)
|
||||
grid_size = int(lut_size ** (1/3))
|
||||
lut_3d = np.zeros((lut_size, lut_size, lut_size, 3), dtype=np.float32)
|
||||
|
||||
for r in range(lut_size):
|
||||
for g in range(lut_size):
|
||||
for b in range(lut_size):
|
||||
# Calculate position in grid
|
||||
grid_x = (b % grid_size) * lut_size + g
|
||||
grid_y = (b // grid_size) * lut_size + r
|
||||
|
||||
if grid_x < width and grid_y < height:
|
||||
lut_3d[r, g, b] = lut_array[grid_y, grid_x]
|
||||
else:
|
||||
# Strip format
|
||||
lut_3d = lut_array.reshape(lut_size, lut_size, lut_size, 3)
|
||||
|
||||
print(f"Successfully loaded image LUT: {lut_size}x{lut_size}x{lut_size}")
|
||||
return torch.from_numpy(lut_3d)
|
||||
|
||||
# If no standard format matches, try to auto-detect
|
||||
total_pixels = width * height
|
||||
cube_root = round(total_pixels ** (1/3))
|
||||
|
||||
if cube_root ** 3 == total_pixels:
|
||||
print(f"Auto-detected LUT size: {cube_root}x{cube_root}x{cube_root}")
|
||||
lut_3d = lut_array.reshape(cube_root, cube_root, cube_root, 3)
|
||||
return torch.from_numpy(lut_3d)
|
||||
|
||||
print(f"Could not determine LUT format for {width}x{height} image")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading image LUT: {e}")
|
||||
return None
|
||||
|
||||
def _apply_lut_tensor(self, images, lut, strength):
|
||||
"""Apply 3D LUT using tensor operations with trilinear interpolation"""
|
||||
batch_size, height, width, channels = images.shape
|
||||
lut_size = lut.shape[0]
|
||||
|
||||
# Scale input to LUT coordinate space
|
||||
coords = images * (lut_size - 1)
|
||||
|
||||
# Get integer and fractional parts
|
||||
coords_floor = torch.floor(coords).long()
|
||||
coords_frac = coords - coords_floor.float()
|
||||
|
||||
# Clamp coordinates
|
||||
coords_floor = torch.clamp(coords_floor, 0, lut_size - 2)
|
||||
coords_ceil = torch.clamp(coords_floor + 1, 0, lut_size - 1)
|
||||
|
||||
# Trilinear interpolation
|
||||
# Get 8 corner values
|
||||
c000 = lut[coords_floor[:, :, :, 0], coords_floor[:, :, :, 1], coords_floor[:, :, :, 2]]
|
||||
c001 = lut[coords_floor[:, :, :, 0], coords_floor[:, :, :, 1], coords_ceil[:, :, :, 2]]
|
||||
c010 = lut[coords_floor[:, :, :, 0], coords_ceil[:, :, :, 1], coords_floor[:, :, :, 2]]
|
||||
c011 = lut[coords_floor[:, :, :, 0], coords_ceil[:, :, :, 1], coords_ceil[:, :, :, 2]]
|
||||
c100 = lut[coords_ceil[:, :, :, 0], coords_floor[:, :, :, 1], coords_floor[:, :, :, 2]]
|
||||
c101 = lut[coords_ceil[:, :, :, 0], coords_floor[:, :, :, 1], coords_ceil[:, :, :, 2]]
|
||||
c110 = lut[coords_ceil[:, :, :, 0], coords_ceil[:, :, :, 1], coords_floor[:, :, :, 2]]
|
||||
c111 = lut[coords_ceil[:, :, :, 0], coords_ceil[:, :, :, 1], coords_ceil[:, :, :, 2]]
|
||||
|
||||
# Interpolate
|
||||
xf, yf, zf = coords_frac[:, :, :, 0:1], coords_frac[:, :, :, 1:2], coords_frac[:, :, :, 2:3]
|
||||
|
||||
c00 = c000 * (1 - xf) + c100 * xf
|
||||
c01 = c001 * (1 - xf) + c101 * xf
|
||||
c10 = c010 * (1 - xf) + c110 * xf
|
||||
c11 = c011 * (1 - xf) + c111 * xf
|
||||
|
||||
c0 = c00 * (1 - yf) + c10 * yf
|
||||
c1 = c01 * (1 - yf) + c11 * yf
|
||||
|
||||
result_lut = c0 * (1 - zf) + c1 * zf
|
||||
|
||||
# Blend with original based on strength
|
||||
result = images * (1 - strength) + result_lut * strength
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
@@ -0,0 +1,174 @@
|
||||
# APNext Cross Processing Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextCrossProcessing:
|
||||
"""
|
||||
APNext Cross Processing Effect Node
|
||||
Simulates cross processing film techniques with channel manipulation and curves
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"process_type": ([
|
||||
"E6_in_C41", "C41_in_E6", "Custom",
|
||||
"Vintage_Warm", "Cold_Blue", "Green_Magenta"
|
||||
], {"default": "E6_in_C41"}),
|
||||
"intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"contrast": ("FLOAT", {"default": 1.2, "min": 0.5, "max": 3.0, "step": 0.01}),
|
||||
"saturation": ("FLOAT", {"default": 1.3, "min": 0.0, "max": 3.0, "step": 0.01}),
|
||||
# Custom channel mixing
|
||||
"red_in_red": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"green_in_red": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"blue_in_red": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"red_in_green": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"green_in_green": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"blue_in_green": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"red_in_blue": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"green_in_blue": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
"blue_in_blue": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_cross_processing"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_cross_processing(self, images, process_type, intensity, contrast, saturation,
|
||||
red_in_red, green_in_red, blue_in_red,
|
||||
red_in_green, green_in_green, blue_in_green,
|
||||
red_in_blue, green_in_blue, blue_in_blue):
|
||||
"""Apply cross processing effect to images"""
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Get the appropriate color matrix
|
||||
if process_type == "Custom":
|
||||
color_matrix = torch.tensor([
|
||||
[red_in_red, green_in_red, blue_in_red],
|
||||
[red_in_green, green_in_green, blue_in_green],
|
||||
[red_in_blue, green_in_blue, blue_in_blue]
|
||||
], device=device, dtype=dtype)
|
||||
else:
|
||||
color_matrix = self._get_preset_matrix(process_type, device, dtype)
|
||||
|
||||
# Apply cross processing
|
||||
result = self._apply_cross_process_tensor(
|
||||
images, color_matrix, intensity, contrast, saturation
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _get_preset_matrix(self, process_type, device, dtype):
|
||||
"""Get predefined color matrices for different cross processing types"""
|
||||
matrices = {
|
||||
"E6_in_C41": torch.tensor([ # Slide film in negative chemistry
|
||||
[1.2, -0.1, 0.1],
|
||||
[0.0, 1.1, -0.1],
|
||||
[-0.1, 0.2, 1.3]
|
||||
], device=device, dtype=dtype),
|
||||
|
||||
"C41_in_E6": torch.tensor([ # Negative film in slide chemistry
|
||||
[0.9, 0.1, -0.1],
|
||||
[0.1, 1.0, 0.1],
|
||||
[0.2, -0.2, 0.8]
|
||||
], device=device, dtype=dtype),
|
||||
|
||||
"Vintage_Warm": torch.tensor([
|
||||
[1.3, 0.1, -0.1],
|
||||
[-0.1, 1.1, 0.2],
|
||||
[-0.2, -0.1, 0.9]
|
||||
], device=device, dtype=dtype),
|
||||
|
||||
"Cold_Blue": torch.tensor([
|
||||
[0.8, -0.1, 0.2],
|
||||
[0.1, 1.0, -0.1],
|
||||
[0.3, 0.2, 1.4]
|
||||
], device=device, dtype=dtype),
|
||||
|
||||
"Green_Magenta": torch.tensor([
|
||||
[1.1, -0.2, 0.1],
|
||||
[-0.3, 1.4, -0.1],
|
||||
[0.1, -0.2, 1.0]
|
||||
], device=device, dtype=dtype),
|
||||
}
|
||||
|
||||
return matrices.get(process_type, torch.eye(3, device=device, dtype=dtype))
|
||||
|
||||
def _apply_cross_process_tensor(self, images, color_matrix, intensity, contrast, saturation):
|
||||
"""Apply cross processing using tensor operations"""
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Apply color matrix transformation
|
||||
images_flat = images.view(-1, 3)
|
||||
processed_flat = torch.matmul(images_flat, color_matrix.t())
|
||||
processed = processed_flat.view(batch_size, height, width, channels)
|
||||
|
||||
# Apply contrast adjustment
|
||||
if contrast != 1.0:
|
||||
mid_point = 0.5
|
||||
processed = (processed - mid_point) * contrast + mid_point
|
||||
|
||||
# Apply saturation adjustment
|
||||
if saturation != 1.0:
|
||||
processed = self._apply_saturation_tensor(processed, saturation)
|
||||
|
||||
# Apply cross processing curves (S-curves for each channel)
|
||||
processed = self._apply_film_curves_tensor(processed)
|
||||
|
||||
# Blend with original based on intensity
|
||||
result = images * (1 - intensity) + processed * intensity
|
||||
|
||||
# Clamp values
|
||||
result = torch.clamp(result, 0, 1)
|
||||
return result
|
||||
|
||||
def _apply_saturation_tensor(self, images, saturation):
|
||||
"""Apply saturation adjustment using tensor operations"""
|
||||
# Calculate luminance
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Blend between grayscale and original based on saturation
|
||||
result = luminance + (images - luminance) * saturation
|
||||
return result
|
||||
|
||||
def _apply_film_curves_tensor(self, images):
|
||||
"""Apply characteristic film curves that create the cross-processing look"""
|
||||
# Create S-curve for each channel with different characteristics
|
||||
# Red channel: lifted shadows, compressed highlights
|
||||
red = images[:, :, :, 0:1]
|
||||
red_curved = self._s_curve_tensor(red, shadows=0.1, highlights=-0.1, contrast=1.1)
|
||||
|
||||
# Green channel: more aggressive S-curve
|
||||
green = images[:, :, :, 1:2]
|
||||
green_curved = self._s_curve_tensor(green, shadows=0.05, highlights=-0.05, contrast=1.2)
|
||||
|
||||
# Blue channel: inverted curve characteristics
|
||||
blue = images[:, :, :, 2:3]
|
||||
blue_curved = self._s_curve_tensor(blue, shadows=-0.05, highlights=0.15, contrast=1.1)
|
||||
|
||||
return torch.cat([red_curved, green_curved, blue_curved], dim=3)
|
||||
|
||||
def _s_curve_tensor(self, channel, shadows=0.0, highlights=0.0, contrast=1.0):
|
||||
"""Apply S-curve to a single channel"""
|
||||
# Apply shadow/highlight adjustments
|
||||
adjusted = channel + shadows * (1 - channel) + highlights * channel
|
||||
|
||||
# Apply contrast with S-curve
|
||||
# Use a sigmoid-based S-curve
|
||||
mid_point = 0.5
|
||||
normalized = (adjusted - mid_point) * contrast
|
||||
s_curved = torch.sigmoid(normalized * 2)
|
||||
|
||||
# Scale back to 0-1 range
|
||||
result = s_curved * (1 - mid_point) + mid_point
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,205 @@
|
||||
# APNext Film Halation Effect Node
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextFilmHalation:
|
||||
"""
|
||||
APNext Film Halation Effect Node
|
||||
Simulates film halation - light bleeding around bright objects
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"radius": ("FLOAT", {"default": 20.0, "min": 1.0, "max": 100.0, "step": 0.5}),
|
||||
"color_tint": (["none", "warm", "cool", "red", "blue", "custom"], {"default": "warm"}),
|
||||
"custom_color_r": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"custom_color_g": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"custom_color_b": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"falloff": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 4.0, "step": 0.1}),
|
||||
"film_type": (["color_negative", "slide_film", "vintage", "modern"], {"default": "color_negative"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_film_halation"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_film_halation(self, images, intensity, threshold, radius, color_tint,
|
||||
custom_color_r, custom_color_g, custom_color_b, falloff, film_type):
|
||||
"""Apply film halation effect to images"""
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Apply halation effect
|
||||
result = self._create_halation_effect(
|
||||
images, intensity, threshold, radius, color_tint,
|
||||
custom_color_r, custom_color_g, custom_color_b, falloff, film_type
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _create_halation_effect(self, images, intensity, threshold, radius, color_tint,
|
||||
custom_color_r, custom_color_g, custom_color_b, falloff, film_type):
|
||||
"""Create film halation effect using tensor operations"""
|
||||
batch_size, height, width, channels = images.shape
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Calculate luminance for threshold detection
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
|
||||
# Create mask for bright areas
|
||||
bright_mask = (luminance > threshold).float()
|
||||
|
||||
# Extract bright areas
|
||||
bright_areas = images * bright_mask.unsqueeze(-1)
|
||||
|
||||
# Apply film-specific characteristics
|
||||
bright_areas = self._apply_film_characteristics(bright_areas, film_type)
|
||||
|
||||
# Create halation glow
|
||||
halation_glow = self._create_glow_effect(bright_areas, radius, falloff)
|
||||
|
||||
# Apply color tinting
|
||||
tinted_glow = self._apply_color_tint(
|
||||
halation_glow, color_tint, custom_color_r, custom_color_g, custom_color_b
|
||||
)
|
||||
|
||||
# Blend halation with original image
|
||||
result = self._blend_halation(images, tinted_glow, intensity)
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _apply_film_characteristics(self, bright_areas, film_type):
|
||||
"""Apply film-specific characteristics to bright areas"""
|
||||
if film_type == "color_negative":
|
||||
# Color negative film has strong red/orange halation
|
||||
bright_areas[:, :, :, 0] *= 1.3 # Enhance red
|
||||
bright_areas[:, :, :, 1] *= 1.1 # Slightly enhance green
|
||||
bright_areas[:, :, :, 2] *= 0.8 # Reduce blue
|
||||
|
||||
elif film_type == "slide_film":
|
||||
# Slide film has more neutral halation with slight magenta cast
|
||||
bright_areas[:, :, :, 0] *= 1.1 # Slight red enhancement
|
||||
bright_areas[:, :, :, 1] *= 0.95 # Slight green reduction
|
||||
bright_areas[:, :, :, 2] *= 1.05 # Slight blue enhancement
|
||||
|
||||
elif film_type == "vintage":
|
||||
# Vintage film has warm, yellowish halation
|
||||
bright_areas[:, :, :, 0] *= 1.4 # Strong red
|
||||
bright_areas[:, :, :, 1] *= 1.2 # Enhanced green (yellow)
|
||||
bright_areas[:, :, :, 2] *= 0.6 # Reduced blue
|
||||
|
||||
elif film_type == "modern":
|
||||
# Modern film has more controlled, neutral halation
|
||||
bright_areas *= 1.1 # Slight overall enhancement
|
||||
|
||||
return torch.clamp(bright_areas, 0, 2) # Allow overexposure
|
||||
|
||||
def _create_glow_effect(self, bright_areas, radius, falloff):
|
||||
"""Create glow effect using multiple gaussian blurs"""
|
||||
# Create multiple blur layers for realistic halation
|
||||
glow_layers = []
|
||||
|
||||
# Primary glow (tight)
|
||||
primary_glow = self._gaussian_blur_tensor(bright_areas, radius * 0.3)
|
||||
glow_layers.append(primary_glow * 0.6)
|
||||
|
||||
# Secondary glow (medium)
|
||||
secondary_glow = self._gaussian_blur_tensor(bright_areas, radius * 0.7)
|
||||
glow_layers.append(secondary_glow * 0.3)
|
||||
|
||||
# Tertiary glow (wide)
|
||||
tertiary_glow = self._gaussian_blur_tensor(bright_areas, radius)
|
||||
glow_layers.append(tertiary_glow * 0.1)
|
||||
|
||||
# Combine glow layers
|
||||
combined_glow = sum(glow_layers)
|
||||
|
||||
# Apply falloff
|
||||
if falloff != 1.0:
|
||||
# Create distance-based falloff
|
||||
luminance = 0.299 * combined_glow[:, :, :, 0] + 0.587 * combined_glow[:, :, :, 1] + 0.114 * combined_glow[:, :, :, 2]
|
||||
falloff_factor = torch.pow(luminance.unsqueeze(-1), 1.0 / falloff)
|
||||
combined_glow = combined_glow * falloff_factor
|
||||
|
||||
return combined_glow
|
||||
|
||||
def _apply_color_tint(self, glow, color_tint, custom_r, custom_g, custom_b):
|
||||
"""Apply color tinting to the halation glow"""
|
||||
if color_tint == "none":
|
||||
return glow
|
||||
|
||||
tint_colors = {
|
||||
"warm": torch.tensor([1.2, 1.0, 0.8], device=glow.device, dtype=glow.dtype),
|
||||
"cool": torch.tensor([0.8, 1.0, 1.2], device=glow.device, dtype=glow.dtype),
|
||||
"red": torch.tensor([1.5, 0.8, 0.8], device=glow.device, dtype=glow.dtype),
|
||||
"blue": torch.tensor([0.8, 0.9, 1.4], device=glow.device, dtype=glow.dtype),
|
||||
"custom": torch.tensor([custom_r, custom_g, custom_b], device=glow.device, dtype=glow.dtype)
|
||||
}
|
||||
|
||||
if color_tint in tint_colors:
|
||||
tint = tint_colors[color_tint].view(1, 1, 1, 3)
|
||||
tinted_glow = glow * tint
|
||||
else:
|
||||
tinted_glow = glow
|
||||
|
||||
return tinted_glow
|
||||
|
||||
def _blend_halation(self, original, halation, intensity):
|
||||
"""Blend halation effect with original image"""
|
||||
# Use screen blend mode for realistic halation
|
||||
# Screen formula: 1 - (1 - base) * (1 - overlay)
|
||||
halation_scaled = halation * intensity
|
||||
|
||||
# Apply screen blending
|
||||
result = 1 - (1 - original) * (1 - halation_scaled)
|
||||
|
||||
return result
|
||||
|
||||
def _gaussian_blur_tensor(self, images, radius):
|
||||
"""Apply gaussian blur using torch operations"""
|
||||
if radius < 0.5:
|
||||
return images
|
||||
|
||||
# Convert radius to sigma
|
||||
sigma = radius / 3.0
|
||||
|
||||
# Create gaussian kernel
|
||||
kernel_size = int(2 * radius + 1)
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
|
||||
# Limit kernel size for performance
|
||||
kernel_size = min(kernel_size, 99)
|
||||
|
||||
# Create 1D gaussian kernel
|
||||
x = torch.arange(kernel_size, dtype=images.dtype, device=images.device)
|
||||
x = x - kernel_size // 2
|
||||
kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2)
|
||||
kernel_1d = kernel_1d / kernel_1d.sum()
|
||||
|
||||
# Apply separable gaussian blur
|
||||
# Reshape for conv operations: [batch, channels, height, width]
|
||||
images_conv = images.permute(0, 3, 1, 2)
|
||||
|
||||
# Horizontal blur
|
||||
kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1)
|
||||
blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3)
|
||||
|
||||
# Vertical blur
|
||||
kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1)
|
||||
blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3)
|
||||
|
||||
# Reshape back to original format: [batch, height, width, channels]
|
||||
return blurred.permute(0, 2, 3, 1)
|
||||
@@ -0,0 +1,315 @@
|
||||
# APNext Glitch Art Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import random
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextGlitchArt:
|
||||
"""
|
||||
APNext Glitch Art Effect Node
|
||||
Creates digital corruption and glitch effects
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"glitch_type": ([
|
||||
"RGB_Shift", "Data_Moshing", "Pixel_Sort", "Scanlines",
|
||||
"Digital_Noise", "Color_Channel_Shift", "Compression_Artifacts"
|
||||
], {"default": "RGB_Shift"}),
|
||||
"intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"randomness": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_glitch_art"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_glitch_art(self, images, glitch_type, intensity, randomness, seed):
|
||||
"""Apply glitch art effects to images"""
|
||||
# Set random seed
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed % (2**32))
|
||||
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
if glitch_type == "RGB_Shift":
|
||||
result = self._rgb_channel_shift(images, intensity, randomness)
|
||||
elif glitch_type == "Data_Moshing":
|
||||
result = self._data_moshing_effect(images, intensity, randomness)
|
||||
elif glitch_type == "Pixel_Sort":
|
||||
result = self._pixel_sorting_effect(images, intensity, randomness)
|
||||
elif glitch_type == "Scanlines":
|
||||
result = self._scanlines_effect(images, intensity, randomness)
|
||||
elif glitch_type == "Digital_Noise":
|
||||
result = self._digital_noise_effect(images, intensity, randomness)
|
||||
elif glitch_type == "Color_Channel_Shift":
|
||||
result = self._color_channel_shift(images, intensity, randomness)
|
||||
elif glitch_type == "Compression_Artifacts":
|
||||
result = self._compression_artifacts(images, intensity, randomness)
|
||||
else:
|
||||
result = images
|
||||
|
||||
return (result,)
|
||||
|
||||
def _rgb_channel_shift(self, images, intensity, randomness):
|
||||
"""Create RGB channel shift glitch effect"""
|
||||
batch_size, height, width, channels = images.shape
|
||||
device = images.device
|
||||
|
||||
# Calculate shift amounts
|
||||
max_shift = int(intensity * 20 + randomness * 10)
|
||||
|
||||
# Create random shifts for each channel
|
||||
r_shift_x = random.randint(-max_shift, max_shift)
|
||||
r_shift_y = random.randint(-max_shift, max_shift)
|
||||
g_shift_x = random.randint(-max_shift, max_shift)
|
||||
g_shift_y = random.randint(-max_shift, max_shift)
|
||||
b_shift_x = random.randint(-max_shift, max_shift)
|
||||
b_shift_y = random.randint(-max_shift, max_shift)
|
||||
|
||||
result = images.clone()
|
||||
|
||||
# Apply shifts to each channel
|
||||
for batch_idx in range(batch_size):
|
||||
# Red channel shift
|
||||
if abs(r_shift_x) > 0 or abs(r_shift_y) > 0:
|
||||
shifted_red = torch.roll(images[batch_idx, :, :, 0], shifts=(r_shift_y, r_shift_x), dims=(0, 1))
|
||||
result[batch_idx, :, :, 0] = shifted_red
|
||||
|
||||
# Green channel shift
|
||||
if abs(g_shift_x) > 0 or abs(g_shift_y) > 0:
|
||||
shifted_green = torch.roll(images[batch_idx, :, :, 1], shifts=(g_shift_y, g_shift_x), dims=(0, 1))
|
||||
result[batch_idx, :, :, 1] = shifted_green
|
||||
|
||||
# Blue channel shift
|
||||
if abs(b_shift_x) > 0 or abs(b_shift_y) > 0:
|
||||
shifted_blue = torch.roll(images[batch_idx, :, :, 2], shifts=(b_shift_y, b_shift_x), dims=(0, 1))
|
||||
result[batch_idx, :, :, 2] = shifted_blue
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _data_moshing_effect(self, images, intensity, randomness):
|
||||
"""Create data moshing effect by corrupting image data"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Number of corruption blocks
|
||||
num_blocks = int(intensity * 50 + randomness * 20)
|
||||
|
||||
for batch_idx in range(batch_size):
|
||||
for _ in range(num_blocks):
|
||||
# Random block position and size
|
||||
block_size = random.randint(5, int(intensity * 50 + 10))
|
||||
x = random.randint(0, width - block_size)
|
||||
y = random.randint(0, height - block_size)
|
||||
|
||||
# Corruption type
|
||||
corruption_type = random.choice(['duplicate', 'scramble', 'invert', 'noise'])
|
||||
|
||||
if corruption_type == 'duplicate':
|
||||
# Duplicate nearby block
|
||||
source_x = max(0, min(width - block_size, x + random.randint(-30, 30)))
|
||||
source_y = max(0, min(height - block_size, y + random.randint(-30, 30)))
|
||||
result[batch_idx, y:y+block_size, x:x+block_size] = \
|
||||
images[batch_idx, source_y:source_y+block_size, source_x:source_x+block_size]
|
||||
|
||||
elif corruption_type == 'scramble':
|
||||
# Scramble pixels in block
|
||||
block = result[batch_idx, y:y+block_size, x:x+block_size].clone()
|
||||
flat_block = block.view(-1, channels)
|
||||
indices = torch.randperm(flat_block.size(0))
|
||||
scrambled = flat_block[indices].view(block_size, block_size, channels)
|
||||
result[batch_idx, y:y+block_size, x:x+block_size] = scrambled
|
||||
|
||||
elif corruption_type == 'invert':
|
||||
# Invert colors in block
|
||||
result[batch_idx, y:y+block_size, x:x+block_size] = \
|
||||
1.0 - result[batch_idx, y:y+block_size, x:x+block_size]
|
||||
|
||||
elif corruption_type == 'noise':
|
||||
# Add noise to block
|
||||
noise = torch.randn_like(result[batch_idx, y:y+block_size, x:x+block_size]) * 0.1
|
||||
result[batch_idx, y:y+block_size, x:x+block_size] += noise
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _pixel_sorting_effect(self, images, intensity, randomness):
|
||||
"""Create pixel sorting glitch effect"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Number of sorting operations
|
||||
num_sorts = int(intensity * 20 + randomness * 10)
|
||||
|
||||
for batch_idx in range(batch_size):
|
||||
for _ in range(num_sorts):
|
||||
# Random sorting direction and area
|
||||
if random.random() < 0.5:
|
||||
# Horizontal sorting
|
||||
row = random.randint(0, height - 1)
|
||||
start_col = random.randint(0, width - 20)
|
||||
end_col = min(width, start_col + random.randint(10, int(intensity * 100 + 20)))
|
||||
|
||||
# Sort pixels by brightness
|
||||
row_data = result[batch_idx, row, start_col:end_col]
|
||||
brightness = 0.299 * row_data[:, 0] + 0.587 * row_data[:, 1] + 0.114 * row_data[:, 2]
|
||||
sorted_indices = torch.argsort(brightness)
|
||||
result[batch_idx, row, start_col:end_col] = row_data[sorted_indices]
|
||||
else:
|
||||
# Vertical sorting
|
||||
col = random.randint(0, width - 1)
|
||||
start_row = random.randint(0, height - 20)
|
||||
end_row = min(height, start_row + random.randint(10, int(intensity * 100 + 20)))
|
||||
|
||||
# Sort pixels by brightness
|
||||
col_data = result[batch_idx, start_row:end_row, col]
|
||||
brightness = 0.299 * col_data[:, 0] + 0.587 * col_data[:, 1] + 0.114 * col_data[:, 2]
|
||||
sorted_indices = torch.argsort(brightness)
|
||||
result[batch_idx, start_row:end_row, col] = col_data[sorted_indices]
|
||||
|
||||
return result
|
||||
|
||||
def _scanlines_effect(self, images, intensity, randomness):
|
||||
"""Create scanlines and CRT-like glitch effects"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Scanline parameters
|
||||
line_spacing = max(2, int(10 - intensity * 8))
|
||||
line_intensity = intensity * 0.5
|
||||
|
||||
for batch_idx in range(batch_size):
|
||||
# Add horizontal scanlines
|
||||
for y in range(0, height, line_spacing):
|
||||
if random.random() < 0.7: # Not every line
|
||||
# Darken scanline
|
||||
result[batch_idx, y] *= (1.0 - line_intensity)
|
||||
|
||||
# Add some color shift
|
||||
if random.random() < randomness:
|
||||
shift = random.randint(-2, 2)
|
||||
if shift != 0:
|
||||
result[batch_idx, y] = torch.roll(result[batch_idx, y], shift, dims=0)
|
||||
|
||||
# Add random glitch lines
|
||||
num_glitch_lines = int(intensity * 5 + randomness * 3)
|
||||
for _ in range(num_glitch_lines):
|
||||
y = random.randint(0, height - 1)
|
||||
# Duplicate or shift line
|
||||
if random.random() < 0.5:
|
||||
# Duplicate from nearby line
|
||||
source_y = max(0, min(height - 1, y + random.randint(-5, 5)))
|
||||
result[batch_idx, y] = result[batch_idx, source_y]
|
||||
else:
|
||||
# Color shift
|
||||
result[batch_idx, y, :, 0] = torch.roll(result[batch_idx, y, :, 0], random.randint(-10, 10))
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _digital_noise_effect(self, images, intensity, randomness):
|
||||
"""Add digital noise and artifacts"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Add random digital noise
|
||||
noise_amount = intensity * randomness * 0.1
|
||||
noise = torch.randn_like(result) * noise_amount
|
||||
result += noise
|
||||
|
||||
# Add salt and pepper noise
|
||||
salt_pepper_prob = intensity * randomness * 0.05
|
||||
mask = torch.rand_like(result[:, :, :, 0:1]) < salt_pepper_prob
|
||||
salt_pepper = torch.rand_like(result) > 0.5
|
||||
result = torch.where(mask, salt_pepper.float(), result)
|
||||
|
||||
# Add compression-like artifacts
|
||||
if intensity > 0.3:
|
||||
# Quantize colors
|
||||
quantization_levels = max(8, int(256 - intensity * 200))
|
||||
result = torch.round(result * quantization_levels) / quantization_levels
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _color_channel_shift(self, images, intensity, randomness):
|
||||
"""Shift color channels in different directions"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Create different shift patterns for each channel
|
||||
for batch_idx in range(batch_size):
|
||||
# Red channel - horizontal shift
|
||||
h_shift = int(intensity * randomness * 20)
|
||||
if h_shift > 0:
|
||||
result[batch_idx, :, :, 0] = torch.roll(result[batch_idx, :, :, 0], h_shift, dims=1)
|
||||
|
||||
# Green channel - vertical shift
|
||||
v_shift = int(intensity * randomness * 15)
|
||||
if v_shift > 0:
|
||||
result[batch_idx, :, :, 1] = torch.roll(result[batch_idx, :, :, 1], v_shift, dims=0)
|
||||
|
||||
# Blue channel - diagonal effect
|
||||
if intensity > 0.5:
|
||||
# Create diagonal shift effect
|
||||
for y in range(height):
|
||||
shift_amount = int((y / height) * intensity * 10)
|
||||
if shift_amount > 0:
|
||||
result[batch_idx, y, :, 2] = torch.roll(result[batch_idx, y, :, 2], shift_amount)
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _compression_artifacts(self, images, intensity, randomness):
|
||||
"""Simulate compression artifacts and blocking"""
|
||||
result = images.clone()
|
||||
batch_size, height, width, channels = images.shape
|
||||
|
||||
# Block size for compression artifacts
|
||||
block_size = max(4, int(16 - intensity * 10))
|
||||
|
||||
for batch_idx in range(batch_size):
|
||||
# Create blocking artifacts
|
||||
for y in range(0, height, block_size):
|
||||
for x in range(0, width, block_size):
|
||||
end_y = min(height, y + block_size)
|
||||
end_x = min(width, x + block_size)
|
||||
|
||||
if random.random() < intensity * randomness:
|
||||
# Average the block (compression-like effect)
|
||||
block = result[batch_idx, y:end_y, x:end_x]
|
||||
avg_color = torch.mean(block, dim=(0, 1), keepdim=True)
|
||||
result[batch_idx, y:end_y, x:end_x] = avg_color
|
||||
|
||||
elif random.random() < intensity * 0.3:
|
||||
# Add ringing artifacts around edges
|
||||
block = result[batch_idx, y:end_y, x:end_x]
|
||||
# Simple edge detection
|
||||
if block.size(0) > 1 and block.size(1) > 1:
|
||||
# Calculate edges safely with proper size handling
|
||||
vertical_edges = torch.abs(block[1:, :, :] - block[:-1, :, :])
|
||||
horizontal_edges = torch.abs(block[:, 1:, :] - block[:, :-1, :])
|
||||
|
||||
# Pad edges to match original block size
|
||||
vertical_edges_padded = torch.zeros_like(block)
|
||||
horizontal_edges_padded = torch.zeros_like(block)
|
||||
|
||||
vertical_edges_padded[:-1, :, :] = vertical_edges
|
||||
horizontal_edges_padded[:, :-1, :] = horizontal_edges
|
||||
|
||||
edges = vertical_edges_padded + horizontal_edges_padded
|
||||
edge_strength = torch.mean(edges)
|
||||
|
||||
if edge_strength > 0.1:
|
||||
# Add ringing
|
||||
ringing = torch.sin(torch.arange(block.size(0), device=images.device).float() * 3.14159) * 0.05 * intensity
|
||||
ringing = ringing.view(-1, 1, 1).expand_as(block)
|
||||
result[batch_idx, y:end_y, x:end_x] += ringing
|
||||
|
||||
return torch.clamp(result, 0, 1)
|
||||
@@ -0,0 +1,237 @@
|
||||
# APNext HDR Tone Mapping Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextHDRToneMapping:
|
||||
"""
|
||||
APNext HDR Tone Mapping Effect Node
|
||||
Simulates HDR tone mapping effects for dramatic and surreal looks
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"method": ([
|
||||
"Reinhard", "Drago", "Mantiuk", "Photographic",
|
||||
"Adaptive_Log", "Filmic", "ACES"
|
||||
], {"default": "Reinhard"}),
|
||||
"exposure": ("FLOAT", {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.01}),
|
||||
"gamma": ("FLOAT", {"default": 2.2, "min": 0.5, "max": 4.0, "step": 0.01}),
|
||||
"intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}),
|
||||
# Method-specific parameters
|
||||
"white_point": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01}),
|
||||
"key_value": ("FLOAT", {"default": 0.18, "min": 0.01, "max": 1.0, "step": 0.01}),
|
||||
"adaptation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"color_correction": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"local_adaptation": ("BOOLEAN", {"default": False}),
|
||||
"radius": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_hdr_tone_mapping"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_hdr_tone_mapping(self, images, method, exposure, gamma, intensity,
|
||||
white_point, key_value, adaptation, color_correction,
|
||||
local_adaptation, radius):
|
||||
"""Apply HDR tone mapping effect to images"""
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Apply exposure adjustment first
|
||||
if exposure != 0.0:
|
||||
exposure_factor = 2.0 ** exposure
|
||||
images = images * exposure_factor
|
||||
|
||||
# Apply tone mapping based on selected method
|
||||
if method == "Reinhard":
|
||||
result = self._reinhard_tone_mapping(images, white_point, local_adaptation, radius)
|
||||
elif method == "Drago":
|
||||
result = self._drago_tone_mapping(images, adaptation, white_point)
|
||||
elif method == "Mantiuk":
|
||||
result = self._mantiuk_tone_mapping(images, color_correction, adaptation)
|
||||
elif method == "Photographic":
|
||||
result = self._photographic_tone_mapping(images, key_value, white_point)
|
||||
elif method == "Adaptive_Log":
|
||||
result = self._adaptive_log_tone_mapping(images, adaptation)
|
||||
elif method == "Filmic":
|
||||
result = self._filmic_tone_mapping(images, white_point)
|
||||
elif method == "ACES":
|
||||
result = self._aces_tone_mapping(images)
|
||||
else:
|
||||
result = images
|
||||
|
||||
# Apply gamma correction
|
||||
if gamma != 1.0:
|
||||
result = torch.pow(torch.clamp(result, 0.001, 1.0), 1.0 / gamma)
|
||||
|
||||
# Blend with original based on intensity
|
||||
result = images * (1 - intensity) + result * intensity
|
||||
|
||||
# Clamp final result
|
||||
result = torch.clamp(result, 0, 1)
|
||||
return (result,)
|
||||
|
||||
def _reinhard_tone_mapping(self, images, white_point, local_adaptation, radius):
|
||||
"""Reinhard tone mapping operator"""
|
||||
# Calculate luminance
|
||||
luminance = self._calculate_luminance(images)
|
||||
|
||||
if local_adaptation:
|
||||
# Local adaptation version (simplified)
|
||||
local_lum = self._gaussian_blur_luminance(luminance, radius)
|
||||
adaptation_lum = local_lum / (local_lum + 1.0)
|
||||
tone_mapped_lum = luminance / (luminance + adaptation_lum)
|
||||
else:
|
||||
# Global version
|
||||
tone_mapped_lum = luminance / (luminance + white_point ** 2)
|
||||
|
||||
# Apply to color channels
|
||||
result = self._apply_luminance_change(images, luminance, tone_mapped_lum)
|
||||
return result
|
||||
|
||||
def _drago_tone_mapping(self, images, adaptation, white_point):
|
||||
"""Drago tone mapping operator"""
|
||||
luminance = self._calculate_luminance(images)
|
||||
|
||||
# Drago's adaptive logarithmic mapping
|
||||
log_lum = torch.log10(torch.clamp(luminance, 1e-6, float('inf')))
|
||||
log_white = torch.log10(white_point)
|
||||
|
||||
# Bias function
|
||||
bias = torch.pow(adaptation, log_lum / log_white)
|
||||
|
||||
# Apply tone mapping
|
||||
tone_mapped_lum = (log_lum / log_white) / torch.log10(2 + 8 * bias)
|
||||
tone_mapped_lum = torch.clamp(tone_mapped_lum, 0, 1)
|
||||
|
||||
result = self._apply_luminance_change(images, luminance, tone_mapped_lum)
|
||||
return result
|
||||
|
||||
def _mantiuk_tone_mapping(self, images, color_correction, adaptation):
|
||||
"""Mantiuk tone mapping operator (simplified version)"""
|
||||
luminance = self._calculate_luminance(images)
|
||||
|
||||
# Logarithmic compression
|
||||
log_lum = torch.log(torch.clamp(luminance, 1e-6, float('inf')))
|
||||
|
||||
# Adaptive factor
|
||||
adaptive_factor = adaptation * torch.tanh(log_lum)
|
||||
|
||||
# Tone mapping
|
||||
tone_mapped_lum = torch.exp(log_lum * adaptive_factor)
|
||||
tone_mapped_lum = tone_mapped_lum / (tone_mapped_lum + 1.0)
|
||||
|
||||
# Color correction
|
||||
if color_correction != 1.0:
|
||||
saturation_factor = 1.0 + (color_correction - 1.0) * (1.0 - tone_mapped_lum)
|
||||
result = self._apply_luminance_change(images, luminance, tone_mapped_lum)
|
||||
result = self._apply_saturation_tensor(result, saturation_factor.mean().item())
|
||||
else:
|
||||
result = self._apply_luminance_change(images, luminance, tone_mapped_lum)
|
||||
|
||||
return result
|
||||
|
||||
def _photographic_tone_mapping(self, images, key_value, white_point):
|
||||
"""Photographic tone reproduction operator"""
|
||||
luminance = self._calculate_luminance(images)
|
||||
|
||||
# Calculate log-average luminance
|
||||
log_avg = torch.exp(torch.mean(torch.log(torch.clamp(luminance, 1e-6, float('inf')))))
|
||||
|
||||
# Scale luminance
|
||||
scaled_lum = (key_value / log_avg) * luminance
|
||||
|
||||
# Apply tone mapping with white point
|
||||
tone_mapped_lum = scaled_lum * (1.0 + scaled_lum / (white_point ** 2)) / (1.0 + scaled_lum)
|
||||
|
||||
result = self._apply_luminance_change(images, luminance, tone_mapped_lum)
|
||||
return result
|
||||
|
||||
def _adaptive_log_tone_mapping(self, images, adaptation):
|
||||
"""Adaptive logarithmic tone mapping"""
|
||||
luminance = self._calculate_luminance(images)
|
||||
|
||||
# Logarithmic mapping with adaptation
|
||||
log_lum = torch.log(torch.clamp(luminance, 1e-6, float('inf')))
|
||||
max_log = torch.max(log_lum)
|
||||
min_log = torch.min(log_lum)
|
||||
|
||||
# Normalize and apply adaptation
|
||||
normalized = (log_lum - min_log) / (max_log - min_log + 1e-6)
|
||||
adapted = torch.pow(normalized, adaptation)
|
||||
|
||||
result = self._apply_luminance_change(images, luminance, adapted)
|
||||
return result
|
||||
|
||||
def _filmic_tone_mapping(self, images, white_point):
|
||||
"""Filmic tone mapping (Uncharted 2 style)"""
|
||||
def filmic_curve(x):
|
||||
A, B, C, D, E, F = 0.15, 0.50, 0.10, 0.20, 0.02, 0.30
|
||||
return ((x * (A * x + C * B) + D * E) / (x * (A * x + B) + D * F)) - E / F
|
||||
|
||||
# Apply filmic curve
|
||||
tone_mapped = filmic_curve(images * 2.0) / filmic_curve(torch.tensor(white_point))
|
||||
|
||||
return torch.clamp(tone_mapped, 0, 1)
|
||||
|
||||
def _aces_tone_mapping(self, images):
|
||||
"""ACES filmic tone mapping"""
|
||||
# ACES RRT/ODT approximation
|
||||
a, b, c, d, e = 2.51, 0.03, 2.43, 0.59, 0.14
|
||||
|
||||
tone_mapped = (images * (a * images + b)) / (images * (c * images + d) + e)
|
||||
|
||||
return torch.clamp(tone_mapped, 0, 1)
|
||||
|
||||
def _calculate_luminance(self, images):
|
||||
"""Calculate luminance from RGB"""
|
||||
return 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
|
||||
def _apply_luminance_change(self, images, old_luminance, new_luminance):
|
||||
"""Apply luminance change while preserving color ratios"""
|
||||
old_luminance = old_luminance.unsqueeze(-1)
|
||||
new_luminance = new_luminance.unsqueeze(-1)
|
||||
|
||||
# Avoid division by zero
|
||||
ratio = torch.where(old_luminance > 1e-6, new_luminance / old_luminance, torch.ones_like(old_luminance))
|
||||
|
||||
result = images * ratio
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _gaussian_blur_luminance(self, luminance, radius):
|
||||
"""Apply gaussian blur to luminance for local adaptation"""
|
||||
# Simple box blur approximation for performance
|
||||
kernel_size = max(3, int(radius * 20))
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
|
||||
# Create simple averaging kernel
|
||||
kernel = torch.ones(1, 1, kernel_size, kernel_size, device=luminance.device) / (kernel_size ** 2)
|
||||
|
||||
# Add batch and channel dimensions for conv2d
|
||||
lum_for_conv = luminance.unsqueeze(1) # [batch, 1, height, width]
|
||||
|
||||
# Apply convolution
|
||||
padding = kernel_size // 2
|
||||
blurred = torch.nn.functional.conv2d(lum_for_conv, kernel, padding=padding)
|
||||
|
||||
# Remove extra dimension
|
||||
return blurred.squeeze(1)
|
||||
|
||||
def _apply_saturation_tensor(self, images, saturation):
|
||||
"""Apply saturation adjustment using tensor operations"""
|
||||
# Calculate luminance
|
||||
luminance = self._calculate_luminance(images)
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Blend between grayscale and original based on saturation
|
||||
result = luminance + (images - luminance) * saturation
|
||||
return torch.clamp(result, 0, 1)
|
||||
@@ -0,0 +1,257 @@
|
||||
# APNext Noise Effect Node
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageDraw
|
||||
import numpy as np
|
||||
import random
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
from ...utils.image_utils import tensor2pil, pil2tensor
|
||||
|
||||
try:
|
||||
from sklearn.cluster import KMeans
|
||||
SKLEARN_AVAILABLE = True
|
||||
except ImportError:
|
||||
SKLEARN_AVAILABLE = False
|
||||
|
||||
|
||||
class APNextNoise:
|
||||
"""
|
||||
APNext Noise Effect Node
|
||||
Adds noise to images - can be B&W with transparency or colored based on dominant image colors
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
"noise_type": (["monochrome", "colored", "film_grain"], {"default": "monochrome"}),
|
||||
"blend_mode": (["overlay", "multiply", "screen", "soft_light"], {"default": "overlay"}),
|
||||
"grain_size": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}),
|
||||
"color_count": ("INT", {"default": 3, "min": 2, "max": 8, "step": 1}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_noise"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX"
|
||||
|
||||
def apply_noise(self, images, intensity, noise_type, blend_mode, grain_size, color_count, seed):
|
||||
"""Apply noise effect to images"""
|
||||
result_images = []
|
||||
|
||||
# Set random seed for reproducible results
|
||||
random.seed(seed)
|
||||
# For numpy, we need to clamp to 32-bit range
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
for image in images:
|
||||
# Convert tensor to PIL
|
||||
pil_image = tensor2pil(image)
|
||||
|
||||
# Apply noise effect
|
||||
noisy_image = self._create_noise_effect(
|
||||
pil_image, intensity, noise_type, blend_mode, grain_size, color_count
|
||||
)
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensor = pil2tensor(noisy_image)
|
||||
result_images.append(result_tensor)
|
||||
|
||||
# Stack all processed images
|
||||
return (torch.cat(result_images, dim=0),)
|
||||
|
||||
def _create_noise_effect(self, image, intensity, noise_type, blend_mode, grain_size, color_count):
|
||||
"""Create noise effect on a PIL image"""
|
||||
# Convert to RGB if not already
|
||||
if image.mode != 'RGB':
|
||||
image = image.convert('RGB')
|
||||
|
||||
width, height = image.size
|
||||
|
||||
if noise_type == "monochrome":
|
||||
noise_layer = self._create_monochrome_noise(width, height, grain_size)
|
||||
elif noise_type == "colored":
|
||||
dominant_colors = self._extract_dominant_colors(image, color_count)
|
||||
noise_layer = self._create_colored_noise(width, height, dominant_colors, grain_size)
|
||||
elif noise_type == "film_grain":
|
||||
noise_layer = self._create_film_grain(width, height, grain_size)
|
||||
|
||||
# Blend the noise with the original image
|
||||
result = self._blend_noise(image, noise_layer, blend_mode, intensity)
|
||||
|
||||
return result
|
||||
|
||||
def _create_monochrome_noise(self, width, height, grain_size):
|
||||
"""Create black and white noise with transparency"""
|
||||
# Create noise array
|
||||
noise_size = int(max(width, height) / grain_size)
|
||||
noise = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8)
|
||||
|
||||
# Create PIL image from noise
|
||||
noise_img = Image.fromarray(noise, mode='L')
|
||||
|
||||
# Resize to match target dimensions
|
||||
noise_img = noise_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS)
|
||||
|
||||
# Convert to RGBA with transparency based on noise intensity
|
||||
noise_rgba = Image.new('RGBA', (width, height))
|
||||
noise_array = np.array(noise_img)
|
||||
|
||||
# Create RGBA array where alpha is based on the noise value
|
||||
rgba_array = np.zeros((height, width, 4), dtype=np.uint8)
|
||||
rgba_array[:, :, 0] = noise_array # R
|
||||
rgba_array[:, :, 1] = noise_array # G
|
||||
rgba_array[:, :, 2] = noise_array # B
|
||||
rgba_array[:, :, 3] = noise_array # A (transparency based on noise)
|
||||
|
||||
return Image.fromarray(rgba_array, 'RGBA')
|
||||
|
||||
def _create_colored_noise(self, width, height, colors, grain_size):
|
||||
"""Create colored noise using dominant colors from the image"""
|
||||
# Create noise array
|
||||
noise_size = int(max(width, height) / grain_size)
|
||||
|
||||
# Create RGB noise using the dominant colors
|
||||
noise_r = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8)
|
||||
noise_g = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8)
|
||||
noise_b = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8)
|
||||
|
||||
# Map noise values to dominant colors
|
||||
color_indices = np.random.choice(len(colors), (noise_size, noise_size))
|
||||
|
||||
# Create colored noise
|
||||
colored_noise = np.zeros((noise_size, noise_size, 3), dtype=np.uint8)
|
||||
for i, color in enumerate(colors):
|
||||
mask = color_indices == i
|
||||
colored_noise[mask] = color
|
||||
|
||||
# Add some randomness to the colors
|
||||
variation = np.random.randint(-30, 31, (noise_size, noise_size, 3), dtype=np.int16)
|
||||
colored_noise = np.clip(colored_noise.astype(np.int16) + variation, 0, 255).astype(np.uint8)
|
||||
|
||||
# Create PIL image
|
||||
noise_img = Image.fromarray(colored_noise)
|
||||
|
||||
# Resize to match target dimensions
|
||||
noise_img = noise_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS)
|
||||
|
||||
return noise_img
|
||||
|
||||
def _create_film_grain(self, width, height, grain_size):
|
||||
"""Create film grain effect"""
|
||||
# Create multiple layers of noise for more realistic grain
|
||||
noise_size = int(max(width, height) / grain_size)
|
||||
|
||||
# Fine grain
|
||||
fine_grain = np.random.normal(0, 0.3, (noise_size, noise_size))
|
||||
# Coarse grain
|
||||
coarse_grain = np.random.normal(0, 0.1, (noise_size // 2, noise_size // 2))
|
||||
coarse_grain = np.repeat(np.repeat(coarse_grain, 2, axis=0), 2, axis=1)[:noise_size, :noise_size]
|
||||
|
||||
# Combine grains
|
||||
combined_grain = fine_grain + coarse_grain
|
||||
|
||||
# Convert to 0-255 range
|
||||
grain_normalized = ((combined_grain + 1) * 127.5).clip(0, 255).astype(np.uint8)
|
||||
|
||||
# Create RGB grain
|
||||
grain_rgb = np.stack([grain_normalized] * 3, axis=2)
|
||||
|
||||
# Create PIL image
|
||||
grain_img = Image.fromarray(grain_rgb)
|
||||
|
||||
# Resize to match target dimensions
|
||||
grain_img = grain_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS)
|
||||
|
||||
return grain_img
|
||||
|
||||
def _extract_dominant_colors(self, image, n_colors):
|
||||
"""Extract dominant colors from the image using optimized method"""
|
||||
# Convert image to numpy array and reshape for clustering
|
||||
img_array = np.array(image)
|
||||
pixels = img_array.reshape(-1, 3)
|
||||
|
||||
# Sample pixels for performance (use every 20th pixel for speed)
|
||||
sampled_pixels = pixels[::20]
|
||||
|
||||
if SKLEARN_AVAILABLE:
|
||||
# Use K-means clustering if available
|
||||
try:
|
||||
kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
|
||||
kmeans.fit(sampled_pixels)
|
||||
colors = kmeans.cluster_centers_.astype(int)
|
||||
return colors.tolist()
|
||||
except:
|
||||
pass
|
||||
|
||||
# Fallback: use histogram-based color extraction (faster than K-means)
|
||||
return self._extract_colors_histogram(sampled_pixels, n_colors)
|
||||
|
||||
def _extract_colors_histogram(self, pixels, n_colors):
|
||||
"""Extract colors using histogram-based method (faster fallback)"""
|
||||
# Quantize colors to reduce search space
|
||||
quantized = (pixels // 32) * 32 # Reduce to ~8 levels per channel
|
||||
|
||||
# Find unique colors and their counts
|
||||
unique_colors, counts = np.unique(quantized.reshape(-1, quantized.shape[-1]),
|
||||
axis=0, return_counts=True)
|
||||
|
||||
# Get the most frequent colors
|
||||
most_frequent_indices = np.argsort(counts)[-n_colors:]
|
||||
dominant_colors = unique_colors[most_frequent_indices]
|
||||
|
||||
# Add some variation to avoid too uniform colors
|
||||
variation = np.random.randint(-16, 17, dominant_colors.shape)
|
||||
dominant_colors = np.clip(dominant_colors + variation, 0, 255)
|
||||
|
||||
return dominant_colors.tolist()
|
||||
|
||||
def _blend_noise(self, base_image, noise_layer, blend_mode, intensity):
|
||||
"""Blend noise with the base image using specified blend mode"""
|
||||
# Convert to numpy arrays
|
||||
base_array = np.array(base_image, dtype=np.float32) / 255.0
|
||||
|
||||
# Handle noise layer (might be RGBA or RGB)
|
||||
if noise_layer.mode == 'RGBA':
|
||||
noise_array = np.array(noise_layer, dtype=np.float32)
|
||||
noise_rgb = noise_array[:, :, :3] / 255.0
|
||||
alpha = noise_array[:, :, 3] / 255.0
|
||||
else:
|
||||
noise_rgb = np.array(noise_layer, dtype=np.float32) / 255.0
|
||||
alpha = np.ones((noise_rgb.shape[0], noise_rgb.shape[1])) * 0.5
|
||||
|
||||
# Apply intensity to alpha
|
||||
alpha = alpha * intensity
|
||||
|
||||
# Apply blend mode
|
||||
if blend_mode == "overlay":
|
||||
# Overlay blend mode
|
||||
mask = base_array < 0.5
|
||||
blended = np.where(
|
||||
mask,
|
||||
2 * base_array * noise_rgb,
|
||||
1 - 2 * (1 - base_array) * (1 - noise_rgb)
|
||||
)
|
||||
elif blend_mode == "multiply":
|
||||
blended = base_array * noise_rgb
|
||||
elif blend_mode == "screen":
|
||||
blended = 1 - (1 - base_array) * (1 - noise_rgb)
|
||||
elif blend_mode == "soft_light":
|
||||
blended = np.where(
|
||||
noise_rgb <= 0.5,
|
||||
base_array - (1 - 2 * noise_rgb) * base_array * (1 - base_array),
|
||||
base_array + (2 * noise_rgb - 1) * (np.sqrt(base_array) - base_array)
|
||||
)
|
||||
|
||||
# Apply alpha blending
|
||||
alpha_expanded = np.expand_dims(alpha, axis=2)
|
||||
result = base_array * (1 - alpha_expanded) + blended * alpha_expanded
|
||||
|
||||
# Clamp and convert back to 0-255
|
||||
result = np.clip(result, 0, 1) * 255
|
||||
|
||||
return Image.fromarray(result.astype(np.uint8))
|
||||
@@ -0,0 +1,153 @@
|
||||
# APNext Rough Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
from ...utils.image_utils import tensor2pil, pil2tensor
|
||||
|
||||
|
||||
class APNextRough:
|
||||
"""
|
||||
APNext Rough Effect Node
|
||||
Creates a rough, posterized effect by reducing the image to a limited color palette
|
||||
Based on numpy implementation for better performance
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Simple cache for color palettes
|
||||
self.palette_cache = {}
|
||||
self.max_cache_size = 10
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"color_count": ("INT", {"default": 24, "min": 4, "max": 64, "step": 1}),
|
||||
"aa_factor": ("INT", {"default": 4, "min": 1, "max": 8, "step": 1}),
|
||||
"resize_method": (["BICUBIC", "LANCZOS", "BILINEAR", "NEAREST"], {"default": "BICUBIC"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_rough"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX"
|
||||
|
||||
def apply_rough(self, images, color_count, aa_factor, resize_method):
|
||||
"""Apply rough/posterized effect to images"""
|
||||
result_images = []
|
||||
|
||||
for image in images:
|
||||
# Convert tensor to PIL
|
||||
pil_image = tensor2pil(image)
|
||||
|
||||
# Apply rough effect
|
||||
rough_image = self._create_rough_effect(
|
||||
pil_image, color_count, aa_factor, resize_method
|
||||
)
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensor = pil2tensor(rough_image)
|
||||
result_images.append(result_tensor)
|
||||
|
||||
# Stack all processed images
|
||||
return (torch.cat(result_images, dim=0),)
|
||||
|
||||
def _create_rough_effect(self, image, color_count, aa_factor, resize_method):
|
||||
"""Create rough effect on a PIL image using numpy for efficiency"""
|
||||
# Convert to RGB if not already
|
||||
if image.mode != 'RGB':
|
||||
image = image.convert('RGB')
|
||||
|
||||
# Convert PIL to numpy array
|
||||
img = np.asarray(image, dtype='int64')
|
||||
height, width = img.shape[:2]
|
||||
|
||||
# Try to get cached palette
|
||||
cache_key = self._get_palette_cache_key(img, color_count)
|
||||
if cache_key in self.palette_cache:
|
||||
colors = self.palette_cache[cache_key]
|
||||
print("Using cached palette.")
|
||||
else:
|
||||
# Generate new palette
|
||||
colors = self._generate_palette(img, color_count)
|
||||
|
||||
# Cache the palette
|
||||
if len(self.palette_cache) >= self.max_cache_size:
|
||||
# Remove oldest entry
|
||||
oldest_key = next(iter(self.palette_cache))
|
||||
del self.palette_cache[oldest_key]
|
||||
self.palette_cache[cache_key] = colors
|
||||
|
||||
print("Created palette.")
|
||||
print(colors)
|
||||
|
||||
# Compute color distance of each pixel to each palette color (vectorized)
|
||||
colors_array = np.array(colors)
|
||||
img_flat = img.reshape(-1, 3)
|
||||
|
||||
# Vectorized distance computation
|
||||
distances = np.abs(img_flat[:, None, :] - colors_array[None, :, :]).sum(axis=2)
|
||||
closest_indices = np.argmin(distances, axis=1)
|
||||
closest_indices = closest_indices.reshape(height, width)
|
||||
|
||||
# Create new image with anti-aliasing
|
||||
new_img = np.zeros((height * aa_factor, width * aa_factor, 3), dtype=np.uint8)
|
||||
|
||||
# Vectorized assignment of colors (much faster than nested loops)
|
||||
for y in range(height):
|
||||
for x in range(width):
|
||||
col = colors[closest_indices[y, x]]
|
||||
# Draw rectangles using numpy slicing
|
||||
new_img[aa_factor * y : aa_factor * y + aa_factor,
|
||||
aa_factor * x : aa_factor * x + aa_factor, :] = col
|
||||
|
||||
print("Finished rough image.")
|
||||
|
||||
# Convert back to PIL and resize
|
||||
nim = Image.fromarray(new_img)
|
||||
resize_filter = getattr(Image, resize_method)
|
||||
aim = nim.resize((width, height), resize_filter)
|
||||
|
||||
return aim
|
||||
|
||||
def _get_palette_cache_key(self, img, color_count):
|
||||
"""Generate cache key for palette based on image statistics and color count"""
|
||||
# Use image statistics for cache key (much faster than hashing entire image)
|
||||
img_sample = img[::10, ::10] # Sample every 10th pixel
|
||||
stats = (
|
||||
img_sample.mean(),
|
||||
img_sample.std(),
|
||||
img_sample.min(),
|
||||
img_sample.max(),
|
||||
color_count
|
||||
)
|
||||
return hash(stats)
|
||||
|
||||
def _generate_palette(self, img, color_count):
|
||||
"""Generate color palette using optimized algorithm"""
|
||||
height, width = img.shape[:2]
|
||||
|
||||
# Initialize color palette with black and white
|
||||
colors = [(255, 255, 255), (0, 0, 0)]
|
||||
colors_array = np.array(colors)
|
||||
|
||||
for i in range(color_count - 2):
|
||||
# Vectorized computation of distances to all existing colors
|
||||
img_expanded = img.reshape(-1, 3) # Flatten image
|
||||
|
||||
# Compute Manhattan distance to all existing colors at once
|
||||
distances = np.abs(img_expanded[:, None, :] - colors_array[None, :, :]).sum(axis=2)
|
||||
min_distances = distances.min(axis=1)
|
||||
|
||||
# Find pixel with maximum minimum distance (furthest from all existing colors)
|
||||
max_idx = np.argmax(min_distances)
|
||||
new_color = img_expanded[max_idx]
|
||||
|
||||
print(f"Added color: {new_color}")
|
||||
colors.append(tuple(new_color))
|
||||
colors_array = np.array(colors)
|
||||
|
||||
return colors
|
||||
@@ -0,0 +1,147 @@
|
||||
# APNext Sharpen Effect Node
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextSharpen:
|
||||
"""
|
||||
APNext Sharpen Effect Node
|
||||
Sharpens images using unsharp mask technique for professional results
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
||||
"radius": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
||||
"threshold": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}),
|
||||
"method": (["unsharp_mask", "high_pass", "edge_enhance"], {"default": "unsharp_mask"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_sharpen"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX"
|
||||
|
||||
def apply_sharpen(self, images, strength, radius, threshold, method):
|
||||
"""Apply sharpening effect to images using optimized tensor operations"""
|
||||
# Work directly with tensors for better performance
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Process all images in batch
|
||||
result = self._apply_sharpening_method_tensor(
|
||||
images, strength, radius, threshold, method, device, dtype
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _apply_sharpening_method_tensor(self, images, strength, radius, threshold, method, device, dtype):
|
||||
"""Apply the specified sharpening method using tensor operations"""
|
||||
if method == "unsharp_mask":
|
||||
return self._unsharp_mask_tensor(images, strength, radius, threshold)
|
||||
elif method == "high_pass":
|
||||
return self._high_pass_sharpen_tensor(images, strength, radius)
|
||||
elif method == "edge_enhance":
|
||||
return self._edge_enhance_tensor(images, strength)
|
||||
|
||||
return images
|
||||
|
||||
def _unsharp_mask_tensor(self, images, strength, radius, threshold):
|
||||
"""Apply unsharp mask sharpening using tensor operations"""
|
||||
# Create blurred version using gaussian blur
|
||||
blurred = self._gaussian_blur_tensor(images, radius)
|
||||
|
||||
# Calculate the mask (difference between original and blurred)
|
||||
mask = images - blurred
|
||||
|
||||
# Apply threshold if specified
|
||||
if threshold > 0:
|
||||
threshold_norm = threshold / 255.0
|
||||
threshold_mask = torch.abs(mask) > threshold_norm
|
||||
mask = torch.where(threshold_mask, mask, torch.zeros_like(mask))
|
||||
|
||||
# Apply the unsharp mask
|
||||
sharpened = images + (mask * strength)
|
||||
|
||||
# Clamp values to valid range
|
||||
return torch.clamp(sharpened, 0, 1)
|
||||
|
||||
def _high_pass_sharpen_tensor(self, images, strength, radius):
|
||||
"""Apply high-pass filter sharpening using tensor operations"""
|
||||
# Create heavily blurred version for high-pass
|
||||
blurred = self._gaussian_blur_tensor(images, radius * 2)
|
||||
|
||||
# Create high-pass filter (original - blurred + 0.5)
|
||||
high_pass = images - blurred + 0.5
|
||||
high_pass = torch.clamp(high_pass, 0, 1)
|
||||
|
||||
# Blend with original using overlay mode
|
||||
mask = images < 0.5
|
||||
result = torch.where(
|
||||
mask,
|
||||
2 * images * high_pass,
|
||||
1 - 2 * (1 - images) * (1 - high_pass)
|
||||
)
|
||||
|
||||
# Apply strength
|
||||
result = images + (result - images) * strength
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _edge_enhance_tensor(self, images, strength):
|
||||
"""Apply edge enhancement sharpening using tensor operations"""
|
||||
# Create a simple sharpening kernel
|
||||
kernel = torch.tensor([
|
||||
[0, -1, 0],
|
||||
[-1, 5, -1],
|
||||
[0, -1, 0]
|
||||
], dtype=images.dtype, device=images.device).view(1, 1, 3, 3)
|
||||
|
||||
# Expand kernel for all channels
|
||||
kernel = kernel.expand(3, 1, 3, 3)
|
||||
|
||||
# Apply convolution
|
||||
images_conv = images.permute(0, 3, 1, 2) # [batch, channels, height, width]
|
||||
sharpened_conv = F.conv2d(images_conv, kernel, padding=1, groups=3)
|
||||
sharpened = sharpened_conv.permute(0, 2, 3, 1) # [batch, height, width, channels]
|
||||
|
||||
# Blend with original
|
||||
result = images + (sharpened - images) * strength
|
||||
return torch.clamp(result, 0, 1)
|
||||
|
||||
def _gaussian_blur_tensor(self, images, radius):
|
||||
"""Apply gaussian blur using torch operations (shared with bloom)"""
|
||||
# Convert radius to sigma (approximation)
|
||||
sigma = radius / 3.0
|
||||
|
||||
# Create gaussian kernel
|
||||
kernel_size = int(2 * radius + 1)
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
|
||||
# Create 1D gaussian kernel
|
||||
x = torch.arange(kernel_size, dtype=images.dtype, device=images.device)
|
||||
x = x - kernel_size // 2
|
||||
kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2)
|
||||
kernel_1d = kernel_1d / kernel_1d.sum()
|
||||
|
||||
# Apply separable gaussian blur
|
||||
# Reshape for conv operations: [batch, channels, height, width]
|
||||
images_conv = images.permute(0, 3, 1, 2)
|
||||
|
||||
# Horizontal blur
|
||||
kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1)
|
||||
blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3)
|
||||
|
||||
# Vertical blur
|
||||
kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1)
|
||||
blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3)
|
||||
|
||||
# Reshape back to original format: [batch, height, width, channels]
|
||||
return blurred.permute(0, 2, 3, 1)
|
||||
@@ -0,0 +1,283 @@
|
||||
# APNext Split Toning Effect Node
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from ...utils.constants import CUSTOM_CATEGORY
|
||||
|
||||
|
||||
class APNextSplitToning:
|
||||
"""
|
||||
APNext Split Toning Effect Node
|
||||
Applies different colors to highlights and shadows independently
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
# Highlight controls
|
||||
"highlight_hue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
|
||||
"highlight_saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"highlight_luminance": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
# Shadow controls
|
||||
"shadow_hue": ("FLOAT", {"default": 240.0, "min": 0.0, "max": 360.0, "step": 1.0}),
|
||||
"shadow_saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"shadow_luminance": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
# Balance and blending
|
||||
"balance": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
"midtone_contrast": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
||||
"blend_mode": (["color", "soft_light", "overlay", "multiply"], {"default": "color"}),
|
||||
"opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_split_toning"
|
||||
CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced"
|
||||
|
||||
def apply_split_toning(self, images, highlight_hue, highlight_saturation, highlight_luminance,
|
||||
shadow_hue, shadow_saturation, shadow_luminance,
|
||||
balance, midtone_contrast, blend_mode, opacity):
|
||||
"""Apply split toning effect to images"""
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Apply split toning
|
||||
result = self._apply_split_toning_tensor(
|
||||
images, highlight_hue, highlight_saturation, highlight_luminance,
|
||||
shadow_hue, shadow_saturation, shadow_luminance,
|
||||
balance, midtone_contrast, blend_mode, opacity
|
||||
)
|
||||
|
||||
return (result,)
|
||||
|
||||
def _apply_split_toning_tensor(self, images, h_hue, h_sat, h_lum, s_hue, s_sat, s_lum,
|
||||
balance, midtone_contrast, blend_mode, opacity):
|
||||
"""Apply split toning using tensor operations"""
|
||||
batch_size, height, width, channels = images.shape
|
||||
device = images.device
|
||||
dtype = images.dtype
|
||||
|
||||
# Convert to HSV for easier manipulation
|
||||
hsv_images = self._rgb_to_hsv_tensor(images)
|
||||
|
||||
# Extract luminance for masking
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Adjust luminance threshold based on balance
|
||||
adjusted_h_lum = h_lum + balance * 0.2
|
||||
adjusted_s_lum = s_lum - balance * 0.2
|
||||
adjusted_h_lum = torch.clamp(torch.tensor(adjusted_h_lum), 0.0, 1.0)
|
||||
adjusted_s_lum = torch.clamp(torch.tensor(adjusted_s_lum), 0.0, 1.0)
|
||||
|
||||
# Create smooth masks for highlights and shadows
|
||||
highlight_mask = self._create_luminance_mask(luminance, adjusted_h_lum, softness=0.2)
|
||||
shadow_mask = self._create_luminance_mask(luminance, adjusted_s_lum, invert=True, softness=0.2)
|
||||
|
||||
# Create toned colors
|
||||
highlight_color = self._create_toned_color(h_hue, h_sat, device, dtype)
|
||||
shadow_color = self._create_toned_color(s_hue, s_sat, device, dtype)
|
||||
|
||||
# Apply toning based on blend mode
|
||||
if blend_mode == "color":
|
||||
result = self._apply_color_blend(images, highlight_color, shadow_color,
|
||||
highlight_mask, shadow_mask)
|
||||
elif blend_mode == "soft_light":
|
||||
result = self._apply_soft_light_blend(images, highlight_color, shadow_color,
|
||||
highlight_mask, shadow_mask)
|
||||
elif blend_mode == "overlay":
|
||||
result = self._apply_overlay_blend(images, highlight_color, shadow_color,
|
||||
highlight_mask, shadow_mask)
|
||||
elif blend_mode == "multiply":
|
||||
result = self._apply_multiply_blend(images, highlight_color, shadow_color,
|
||||
highlight_mask, shadow_mask)
|
||||
|
||||
# Apply midtone contrast
|
||||
if midtone_contrast != 0.0:
|
||||
result = self._apply_midtone_contrast_tensor(result, midtone_contrast)
|
||||
|
||||
# Blend with original based on opacity
|
||||
result = images * (1 - opacity) + result * opacity
|
||||
|
||||
# Clamp values
|
||||
result = torch.clamp(result, 0, 1)
|
||||
return result
|
||||
|
||||
def _create_luminance_mask(self, luminance, threshold, invert=False, softness=0.1):
|
||||
"""Create smooth luminance-based mask"""
|
||||
if invert:
|
||||
mask = torch.sigmoid((threshold - luminance) / softness)
|
||||
else:
|
||||
mask = torch.sigmoid((luminance - threshold) / softness)
|
||||
return mask
|
||||
|
||||
def _create_toned_color(self, hue, saturation, device, dtype):
|
||||
"""Create a color from hue and saturation"""
|
||||
# Convert hue from degrees to 0-1 range
|
||||
h = (hue % 360.0) / 360.0
|
||||
s = saturation
|
||||
v = 1.0 # Full brightness
|
||||
|
||||
# Convert HSV to RGB
|
||||
rgb = self._hsv_to_rgb_single(h, s, v, device, dtype)
|
||||
return rgb
|
||||
|
||||
def _hsv_to_rgb_single(self, h, s, v, device, dtype):
|
||||
"""Convert single HSV values to RGB"""
|
||||
h = torch.tensor(h, device=device, dtype=dtype)
|
||||
s = torch.tensor(s, device=device, dtype=dtype)
|
||||
v = torch.tensor(v, device=device, dtype=dtype)
|
||||
|
||||
c = v * s
|
||||
x = c * (1 - torch.abs((h * 6) % 2 - 1))
|
||||
m = v - c
|
||||
|
||||
h_i = torch.floor(h * 6).long()
|
||||
|
||||
# Create RGB based on hue sector
|
||||
if h_i == 0:
|
||||
r, g, b = c, x, 0
|
||||
elif h_i == 1:
|
||||
r, g, b = x, c, 0
|
||||
elif h_i == 2:
|
||||
r, g, b = 0, c, x
|
||||
elif h_i == 3:
|
||||
r, g, b = 0, x, c
|
||||
elif h_i == 4:
|
||||
r, g, b = x, 0, c
|
||||
else:
|
||||
r, g, b = c, 0, x
|
||||
|
||||
return torch.stack([r + m, g + m, b + m])
|
||||
|
||||
def _apply_color_blend(self, images, highlight_color, shadow_color, h_mask, s_mask):
|
||||
"""Apply color blend mode for split toning"""
|
||||
# Convert to HSV
|
||||
hsv = self._rgb_to_hsv_tensor(images)
|
||||
|
||||
# Replace hue and saturation while preserving luminance
|
||||
h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images)
|
||||
s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images)
|
||||
|
||||
# Blend colors
|
||||
toned = images.clone()
|
||||
toned = toned * (1 - h_mask) + h_color_expanded * h_mask
|
||||
toned = toned * (1 - s_mask) + s_color_expanded * s_mask
|
||||
|
||||
return toned
|
||||
|
||||
def _apply_soft_light_blend(self, images, highlight_color, shadow_color, h_mask, s_mask):
|
||||
"""Apply soft light blend mode"""
|
||||
h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images)
|
||||
s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images)
|
||||
|
||||
# Soft light formula
|
||||
h_blend = torch.where(
|
||||
h_color_expanded <= 0.5,
|
||||
images - (1 - 2 * h_color_expanded) * images * (1 - images),
|
||||
images + (2 * h_color_expanded - 1) * (torch.sqrt(images) - images)
|
||||
)
|
||||
|
||||
s_blend = torch.where(
|
||||
s_color_expanded <= 0.5,
|
||||
images - (1 - 2 * s_color_expanded) * images * (1 - images),
|
||||
images + (2 * s_color_expanded - 1) * (torch.sqrt(images) - images)
|
||||
)
|
||||
|
||||
result = images * (1 - h_mask) + h_blend * h_mask
|
||||
result = result * (1 - s_mask) + s_blend * s_mask
|
||||
|
||||
return result
|
||||
|
||||
def _apply_overlay_blend(self, images, highlight_color, shadow_color, h_mask, s_mask):
|
||||
"""Apply overlay blend mode"""
|
||||
h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images)
|
||||
s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images)
|
||||
|
||||
# Overlay formula
|
||||
h_blend = torch.where(
|
||||
images <= 0.5,
|
||||
2 * images * h_color_expanded,
|
||||
1 - 2 * (1 - images) * (1 - h_color_expanded)
|
||||
)
|
||||
|
||||
s_blend = torch.where(
|
||||
images <= 0.5,
|
||||
2 * images * s_color_expanded,
|
||||
1 - 2 * (1 - images) * (1 - s_color_expanded)
|
||||
)
|
||||
|
||||
result = images * (1 - h_mask) + h_blend * h_mask
|
||||
result = result * (1 - s_mask) + s_blend * s_mask
|
||||
|
||||
return result
|
||||
|
||||
def _apply_multiply_blend(self, images, highlight_color, shadow_color, h_mask, s_mask):
|
||||
"""Apply multiply blend mode"""
|
||||
h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images)
|
||||
s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images)
|
||||
|
||||
h_blend = images * h_color_expanded
|
||||
s_blend = images * s_color_expanded
|
||||
|
||||
result = images * (1 - h_mask) + h_blend * h_mask
|
||||
result = result * (1 - s_mask) + s_blend * s_mask
|
||||
|
||||
return result
|
||||
|
||||
def _apply_midtone_contrast_tensor(self, images, contrast):
|
||||
"""Apply contrast specifically to midtones"""
|
||||
# Create midtone mask
|
||||
luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2]
|
||||
luminance = luminance.unsqueeze(-1)
|
||||
|
||||
# Gaussian-like curve centered at 0.5
|
||||
midtone_mask = torch.exp(-((luminance - 0.5) ** 2) / (2 * 0.2 ** 2))
|
||||
|
||||
# Apply contrast
|
||||
contrast_factor = 1.0 + contrast
|
||||
mid_point = 0.5
|
||||
contrasted = (images - mid_point) * contrast_factor + mid_point
|
||||
|
||||
# Blend based on midtone mask
|
||||
result = images * (1 - midtone_mask) + contrasted * midtone_mask
|
||||
return result
|
||||
|
||||
def _rgb_to_hsv_tensor(self, rgb):
|
||||
"""Convert RGB to HSV color space"""
|
||||
r, g, b = rgb[:, :, :, 0], rgb[:, :, :, 1], rgb[:, :, :, 2]
|
||||
|
||||
max_val, max_idx = torch.max(rgb, dim=3)
|
||||
min_val, _ = torch.min(rgb, dim=3)
|
||||
|
||||
delta = max_val - min_val
|
||||
|
||||
# Hue calculation
|
||||
hue = torch.zeros_like(max_val)
|
||||
mask = delta != 0
|
||||
|
||||
# Red is max
|
||||
red_max = (max_idx == 0) & mask
|
||||
hue[red_max] = ((g[red_max] - b[red_max]) / delta[red_max]) % 6
|
||||
|
||||
# Green is max
|
||||
green_max = (max_idx == 1) & mask
|
||||
hue[green_max] = (b[green_max] - r[green_max]) / delta[green_max] + 2
|
||||
|
||||
# Blue is max
|
||||
blue_max = (max_idx == 2) & mask
|
||||
hue[blue_max] = (r[blue_max] - g[blue_max]) / delta[blue_max] + 4
|
||||
|
||||
hue = hue / 6.0
|
||||
|
||||
# Saturation
|
||||
saturation = torch.zeros_like(max_val)
|
||||
saturation[max_val != 0] = delta[max_val != 0] / max_val[max_val != 0]
|
||||
|
||||
# Value
|
||||
value = max_val
|
||||
|
||||
return torch.stack([hue, saturation, value], dim=3)
|
||||
Reference in New Issue
Block a user