nodes.py

from PIL import Image
import numpy as np
import os
import tempfile
from types import SimpleNamespace
try:
    from worker import Worker
except Exception as e:
    Worker = None
    IMPORT_ERROR = str(e)
else:
    IMPORT_ERROR = None

class NovaNodes:
    """
    ComfyUI node: Full post-processing chain using Worker from GUI
    All augmentations with tunable parameters.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),

                # EXIF
                "apply_exif_o": ("BOOLEAN", {"default": True}),

                # Noise
                "noise_std_frac": ("FLOAT", {"default": 0.015, "min": 0.0, "max": 0.1, "step": 0.001}),
                "hot_pixel_prob": ("FLOAT", {"default": 1e-6, "min": 0.0, "max": 1e-3, "step": 1e-7}),
                "perturb_mag_frac": ("FLOAT", {"default": 0.008, "min": 0.0, "max": 0.05, "step": 0.001}),

                # CLAHE
                "clahe_clip": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 10.0, "step": 0.1}),
                "clahe_grid": ("INT", {"default": 8, "min": 2, "max": 32, "step": 1}),

                # Fourier
                "apply_fourier_o": ("BOOLEAN", {"default": True}),
                "fourier_strength": ("FLOAT", {"default": 0.9, "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_alpha": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0, "step": 0.1}),
                "fourier_radial_smooth": ("INT", {"default": 5, "min": 0, "max": 50, "step": 1}),
                "fourier_mode": (["auto", "ref", "model"], {"default": "auto"}),

                # Vignette
                "apply_vignette_o": ("BOOLEAN", {"default": True}),
                "vignette_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),

                # Chromatic aberration
                "apply_chromatic_aberration_o": ("BOOLEAN", {"default": True}),
                "ca_shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),

                # Banding
                "apply_banding_o": ("BOOLEAN", {"default": True}),
                "banding_levels": ("INT", {"default": 64, "min": 2, "max": 256, "step": 1}),

                # Motion blur
                "apply_motion_blur_o": ("BOOLEAN", {"default": True}),
                "motion_blur_ksize": ("INT", {"default": 7, "min": 3, "max": 31, "step": 2}),

                # JPEG cycles
                "apply_jpeg_cycles_o": ("BOOLEAN", {"default": True}),
                "jpeg_cycles": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
                "jpeg_quality": ("INT", {"default": 85, "min": 10, "max": 100, "step": 1}),

                # Camera simulation
                "sim_camera": ("BOOLEAN", {"default": False}),
                "enable_bayer": ("BOOLEAN", {"default": True}),
                "iso_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 16.0, "step": 0.1}),
                "read_noise": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 50.0, "step": 0.1}),
            },
            "optional": {
                "ref_image": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("IMAGE", "EXIF")
    FUNCTION = "process"
    CATEGORY = "postprocessing"

    def process(self, image, ref_image=None,
                apply_exif_o=True,
                noise_std_frac=0.015,
                hot_pixel_prob=1e-6,
                perturb_mag_frac=0.008,
                clahe_clip=2.0,
                clahe_grid=8,
                apply_fourier_o=True,
                fourier_strength=0.9,
                fourier_randomness=0.05,
                fourier_phase_perturb=0.08,
                fourier_alpha=1.0,
                fourier_radial_smooth=5,
                fourier_mode="auto",
                apply_vignette_o=True,
                vignette_strength=0.5,
                apply_chromatic_aberration_o=True,
                ca_shift=1.0,
                apply_banding_o=True,
                banding_levels=64,
                apply_motion_blur_o=True,
                motion_blur_ksize=7,
                apply_jpeg_cycles_o=True,
                jpeg_cycles=2,
                jpeg_quality=85,
                sim_camera=False,
                enable_bayer=True,
                iso_scale=1.0,
                read_noise=2.0):

        if Worker is None:
            raise ImportError(f"Could not import Worker module: {IMPORT_ERROR}")

        # Ensure input image is a PIL Image
        if not isinstance(image, Image.Image):
            raise ValueError("Input image must be a PIL Image object")

        # Save input image as temporary file
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_input:
            input_path = tmp_input.name
            image.save(input_path)

        # Prepare reference image if provided
        ref_path = None
        if ref_image is not None:
            if not isinstance(ref_image, Image.Image):
                raise ValueError("Reference image must be a PIL Image object")
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_ref:
                ref_path = tmp_ref.name
                ref_image.save(ref_path)

        # Create output temporary file path
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_output:
            output_path = tmp_output.name

        # Prepare parameters for Worker
        args = SimpleNamespace(
            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,
            strength=fourier_strength,
            randomness=fourier_randomness,
            phase_perturb=fourier_phase_perturb,
            alpha=fourier_alpha,
            radial_smooth=fourier_radial_smooth,
            fft_mode=fourier_mode,
            vignette_strength=vignette_strength,
            chroma_strength=ca_shift,
            banding_strength=1.0 if apply_banding_o else 0.0,
            motion_blur_kernel=motion_blur_ksize,
            jpeg_cycles=jpeg_cycles,
            jpeg_qmin=jpeg_quality,
            jpeg_qmax=jpeg_quality,
            sim_camera=sim_camera,
            no_no_bayer=enable_bayer,
            iso_scale=iso_scale,
            read_noise=read_noise,
            ref=ref_path,
            fft_ref=ref_path,
            seed=None,  # Seed handling can be added if Worker supports it
            cutoff=0.25  # Default value from GUI, adjustable if needed
        )

        # Run Worker
        worker = Worker(input_path, output_path, args)
        worker.run()  # Assuming Worker has a synchronous run() method

        # Load output image
        output_img = Image.open(output_path)

        # Handle EXIF
        new_exif = ""
        if apply_exif_o:
            output_img, new_exif = self._add_fake_exif(output_img)

        # Clean up temporary files
        os.unlink(input_path)
        if ref_path:
            os.unlink(ref_path)
        os.unlink(output_path)

        return (output_img, new_exif)

    def _add_fake_exif(self, img: Image.Image) -> tuple[Image.Image, str]:
        """Insert random but realistic camera EXIF metadata."""
        import random
        import io
        import piexif
        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)",
}