import torch from PIL import Image import numpy as np import os import tempfile from types import SimpleNamespace from typing import Tuple try: from .image_postprocess import process_image except Exception as e: process_image = None IMPORT_ERROR = str(e) else: IMPORT_ERROR = None lut_extensions = ['png','npy','cube'] class NovaNodes: """ ComfyUI node: Full post-processing chain using process_image from image_postprocess All augmentations with tunable parameters. NOTE: Adjusted to match FOOLAI output: - Returns an IMAGE as a single PyTorch tensor shaped (1, H, W, C), dtype=float32, values in [0.0, 1.0]. - Returns EXIF as a STRING (second output slot). Added LUT support: two new node inputs: - lut: STRING path to a LUT file (1D PNG 256x1, .npy, or .cube). Empty string -> disabled. - lut_strength: FLOAT blend strength (0.0..1.0) """ @classmethod def INPUT_TYPES(s): # --- MODIFICATION: Rearranged inputs and updated defaults to match the reference image --- return { "required": { "image": ("IMAGE",), # Parameters (Manual Mode) - Order and defaults match the image "noise_std_frac": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 0.1, "step": 0.001}), "clahe_clip": ("FLOAT", {"default": 2.00, "min": 0.5, "max": 10.0, "step": 0.1}), "clahe_grid": ("INT", {"default": 8, "min": 2, "max": 32, "step": 1}), "fourier_cutoff": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}), "apply_fourier_o": ("BOOLEAN", {"default": True}), "fourier_strength": ("FLOAT", {"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01}), "fourier_randomness": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 0.5, "step": 0.01}), "fourier_phase_perturb": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 0.5, "step": 0.01}), "fourier_radial_smooth": ("INT", {"default": 5, "min": 0, "max": 50, "step": 1}), "fourier_mode": (["auto", "ref", "model"], {"default": "auto"}), "fourier_alpha": ("FLOAT", {"default": 1.00, "min": 0.1, "max": 4.0, "step": 0.1}), "perturb_mag_frac": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.05, "step": 0.001}), "enable_awb": ("BOOLEAN", {"default": True}), "sim_camera": ("BOOLEAN", {"default": True}), # This corresponds to "Enable camera pipeline simulation" "enable_lut": ("BOOLEAN", {"default": True}), "lut": ("STRING", {"default": "X://insert/path/here(.png/.npy/.cube)", "vhs_path_extensions": lut_extensions}), "lut_strength": ("FLOAT", {"default": 1.00, "min": 0.0, "max": 1.0, "step": 0.01}), # Camera simulator options - Order and defaults match the image "enable_bayer": ("BOOLEAN", {"default": True}), "apply_jpeg_cycles_o": ("BOOLEAN", {"default": True}), "jpeg_cycles": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), "jpeg_quality": ("INT", {"default": 88, "min": 10, "max": 100, "step": 1}), "apply_vignette_o": ("BOOLEAN", {"default": True}), "vignette_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}), "apply_chromatic_aberration_o": ("BOOLEAN", {"default": True}), "ca_shift": ("FLOAT", {"default": 1.20, "min": 0.0, "max": 5.0, "step": 0.1}), "iso_scale": ("FLOAT", {"default": 1.00, "min": 0.1, "max": 16.0, "step": 0.1}), "read_noise": ("FLOAT", {"default": 2.00, "min": 0.0, "max": 50.0, "step": 0.1}), "hot_pixel_prob": ("FLOAT", {"default": 1e-7, "min": 0.0, "max": 1e-3, "step": 1e-7}), "apply_banding_o": ("BOOLEAN", {"default": True}), "banding_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), "apply_motion_blur_o": ("BOOLEAN", {"default": True}), "motion_blur_ksize": ("INT", {"default": 1, "min": 1, "max": 31, "step": 2}), # Other options "apply_exif_o": ("BOOLEAN", {"default": True}), }, "optional": { "awb_ref_image": ("IMAGE",), "fft_ref_image": ("IMAGE",), } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("IMAGE", "EXIF") FUNCTION = "process" CATEGORY = "postprocessing" def process(self, image, noise_std_frac=0.02, clahe_clip=2.0, clahe_grid=8, fourier_cutoff=0.25, apply_fourier_o=True, fourier_strength=0.9, fourier_randomness=0.05, fourier_phase_perturb=0.08, fourier_radial_smooth=5, fourier_mode="auto", fourier_alpha=1.0, perturb_mag_frac=0.01, enable_awb=True, sim_camera=True, enable_lut=True, lut="", lut_strength=1.0, enable_bayer=True, apply_jpeg_cycles_o=True, jpeg_cycles=1, jpeg_quality=88, apply_vignette_o=True, vignette_strength=0.35, apply_chromatic_aberration_o=True, ca_shift=1.20, iso_scale=1.0, read_noise=2.0, hot_pixel_prob=1e-7, apply_banding_o=True, banding_strength=0.0, apply_motion_blur_o=True, motion_blur_ksize=1, apply_exif_o=True, awb_ref_image=None, fft_ref_image=None ): if process_image is None: raise ImportError(f"Could not import process_image function: {IMPORT_ERROR}") tmp_files = [] def to_pil_from_any(inp): """Convert a torch tensor / numpy array of many shapes into a PIL RGB Image.""" if isinstance(inp, torch.Tensor): arr = inp.detach().cpu().numpy() else: arr = np.asarray(inp) if arr.ndim == 4 and arr.shape[0] == 1: arr = arr[0] if arr.ndim == 3 and arr.shape[0] in (1, 3): arr = np.transpose(arr, (1, 2, 0)) if arr.ndim == 2: arr = arr[:, :, None] if arr.ndim != 3: raise TypeError(f"Cannot convert array to HWC image, final ndim={arr.ndim}, shape={arr.shape}") if np.issubdtype(arr.dtype, np.floating): if arr.max() <= 1.0: arr = (arr * 255.0).clip(0, 255).astype(np.uint8) else: arr = np.clip(arr, 0, 255).astype(np.uint8) else: arr = arr.astype(np.uint8) if arr.shape[2] == 1: arr = np.repeat(arr, 3, axis=2) return Image.fromarray(arr) try: # ---- Input image -> temporary input file ---- pil_img = to_pil_from_any(image[0]) with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_input: input_path = tmp_input.name pil_img.save(input_path) tmp_files.append(input_path) # ---- AWB reference image if present ---- awb_ref_path = None if awb_ref_image is not None: pil_ref_awb = to_pil_from_any(awb_ref_image[0]) with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_ref_awb: awb_ref_path = tmp_ref_awb.name pil_ref_awb.save(awb_ref_path) tmp_files.append(awb_ref_path) # ---- FFT reference image if present ---- fft_ref_path = None if fft_ref_image is not None: pil_ref_fft = to_pil_from_any(fft_ref_image[0]) with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_ref_fft: fft_ref_path = tmp_ref_fft.name pil_ref_fft.save(fft_ref_path) tmp_files.append(fft_ref_path) # ---- Output path ---- with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_output: output_path = tmp_output.name tmp_files.append(output_path) # Prepare args for process_image args = SimpleNamespace( input=input_path, output=output_path, awb=enable_awb, # Explicit AWB flag ref=awb_ref_path, noise_std=noise_std_frac, hot_pixel_prob=hot_pixel_prob, perturb=perturb_mag_frac, clahe_clip=clahe_clip, tile=clahe_grid, fstrength=fourier_strength if apply_fourier_o else 0.0, randomness=fourier_randomness, phase_perturb=fourier_phase_perturb, fft_alpha=fourier_alpha, radial_smooth=fourier_radial_smooth, fft_mode=fourier_mode, fft_ref=fft_ref_path, vignette_strength=vignette_strength if apply_vignette_o else 0.0, chroma_strength=ca_shift if apply_chromatic_aberration_o else 0.0, banding_strength=banding_strength if apply_banding_o else 0.0, motion_blur_kernel=motion_blur_ksize if apply_motion_blur_o else 1, jpeg_cycles=jpeg_cycles if apply_jpeg_cycles_o else 1, jpeg_qmin=jpeg_quality, jpeg_qmax=96, # As per image range sim_camera=sim_camera, no_no_bayer=not enable_bayer, # FIX: Inverted logic corrected iso_scale=iso_scale, read_noise=read_noise, seed=None, # Seed is not user-configurable in this version cutoff=fourier_cutoff, lut=(lut if enable_lut and lut != "" else None), lut_strength=lut_strength, ) # ---- Run the processing function ---- process_image(input_path, output_path, args) # ---- Load result (force RGB) ---- output_img = Image.open(output_path).convert("RGB") img_out = np.array(output_img) # ---- EXIF insertion (optional) ---- new_exif = "" if apply_exif_o: try: output_img_with_exif, new_exif = self._add_fake_exif(output_img) output_img = output_img_with_exif img_out = np.array(output_img.convert("RGB")) except Exception: new_exif = "" # ---- Convert to FOOLAI-style tensor: (1, H, W, C), float32 in [0,1] ---- img_float = img_out.astype(np.float32) / 255.0 tensor_out = torch.from_numpy(img_float).to(dtype=torch.float32).unsqueeze(0) tensor_out = torch.clamp(tensor_out, 0.0, 1.0) return (tensor_out, new_exif) finally: for p in tmp_files: try: os.unlink(p) except Exception: pass def _add_fake_exif(self, img: Image.Image) -> Tuple[Image.Image, str]: """Insert random but realistic camera EXIF metadata.""" import random import io try: import piexif except Exception: raise exif_dict = { "0th": { piexif.ImageIFD.Make: random.choice(["Canon", "Nikon", "Sony", "Fujifilm", "Olympus", "Leica"]), piexif.ImageIFD.Model: random.choice([ "EOS 5D Mark III", "D850", "Alpha 7R IV", "X-T4", "OM-D E-M1 Mark III", "Q2" ]), piexif.ImageIFD.Software: "Adobe Lightroom", }, "Exif": { piexif.ExifIFD.FNumber: (random.randint(10, 22), 10), piexif.ExifIFD.ExposureTime: (1, random.randint(60, 4000)), piexif.ExifIFD.ISOSpeedRatings: random.choice([100, 200, 400, 800, 1600, 3200]), piexif.ExifIFD.FocalLength: (random.randint(24, 200), 1), }, } exif_bytes = piexif.dump(exif_dict) output = io.BytesIO() img.save(output, format="JPEG", exif=exif_bytes) output.seek(0) return (Image.open(output), str(exif_bytes)) # ------------- # Registration # ------------- NODE_CLASS_MAPPINGS = { "NovaNodes": NovaNodes, } NODE_DISPLAY_NAME_MAPPINGS = { "NovaNodes": "Image Postprocess (NOVA NODES)", }