ComfyUI Custom Node Development

class MyCustomNode:
    """
    A custom ComfyUI node.
    """

    # Unique identifier (must be unique across all nodes)
    @classmethod
    def INPUT_TYPES(cls):
        """Define node inputs."""
        return {
            "required": {
                "image": ("IMAGE",),
                "strength": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 2.0,
                    "step": 0.1
                }),
            },
            "optional": {
                "mask": ("MASK",),
            },
        }

    # Output types (tuple of type names)
    RETURN_TYPES = ("IMAGE",)

    # Output names (tuple of display names)
    RETURN_NAMES = ("processed_image",)

    # Function to call for execution
    FUNCTION = "process"

    # Category in the node menu
    CATEGORY = "Oshun/Image Processing"

    def process(self, image, strength, mask=None):
        """
        Main processing function.

        Args:
            image: Input image tensor [B, H, W, C]
            strength: Processing strength
            mask: Optional mask tensor

        Returns:
            Tuple of outputs matching RETURN_TYPES
        """
        # Process the image
        result = self.apply_effect(image, strength, mask)

        # Return as tuple
        return (result,)

    def apply_effect(self, image, strength, mask):
        """Apply the custom effect."""
        import torch

        # Your processing logic here
        processed = image * strength

        if mask is not None:
            processed = image * (1 - mask) + processed * mask

        return processed


# Node registration
NODE_CLASS_MAPPINGS = {
    "MyCustomNode": MyCustomNode,
}

# Display names
NODE_DISPLAY_NAME_MAPPINGS = {
    "MyCustomNode": "My Custom Effect",
}

custom_nodes/
└── oshun_nodes/
    ├── __init__.py           # Node registration
    ├── nodes/
    │   ├── __init__.py
    │   ├── image_processing.py
    │   ├── text_processing.py
    │   └── integration.py
    ├── utils/
    │   ├── __init__.py
    │   └── helpers.py
    ├── requirements.txt      # Python dependencies
    └── install.py           # Optional installation script

"""Oshun Custom Nodes for ComfyUI."""

from .nodes.image_processing import (
    OshunImageEnhance,
    OshunColorCorrect,
    OshunStyleTransfer,
)
from .nodes.text_processing import (
    OshunPromptEnhancer,
    OshunNegativeGenerator,
)
from .nodes.integration import (
    OshunS3Upload,
    OshunWebhookNotify,
)

# All nodes to register
NODE_CLASS_MAPPINGS = {
    # Image processing
    "OshunImageEnhance": OshunImageEnhance,
    "OshunColorCorrect": OshunColorCorrect,
    "OshunStyleTransfer": OshunStyleTransfer,

    # Text processing
    "OshunPromptEnhancer": OshunPromptEnhancer,
    "OshunNegativeGenerator": OshunNegativeGenerator,

    # Integration
    "OshunS3Upload": OshunS3Upload,
    "OshunWebhookNotify": OshunWebhookNotify,
}

# Display names
NODE_DISPLAY_NAME_MAPPINGS = {
    "OshunImageEnhance": "Oshun Image Enhance",
    "OshunColorCorrect": "Oshun Color Correct",
    "OshunStyleTransfer": "Oshun Style Transfer",
    "OshunPromptEnhancer": "Oshun Prompt Enhancer",
    "OshunNegativeGenerator": "Oshun Negative Generator",
    "OshunS3Upload": "Oshun S3 Upload",
    "OshunWebhookNotify": "Oshun Webhook Notify",
}

# Version
__version__ = "1.0.0"

import torch
import torch.nn.functional as F
from typing import Tuple, Optional

class OshunImageEnhance:
    """
    Enhance image quality with multiple adjustable parameters.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "contrast": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.05,
                    "display": "slider"
                }),
                "brightness": ("FLOAT", {
                    "default": 0.0,
                    "min": -1.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "slider"
                }),
                "saturation": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.05,
                    "display": "slider"
                }),
                "sharpness": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 2.0,
                    "step": 0.05,
                    "display": "slider"
                }),
            },
            "optional": {
                "mask": ("MASK",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("enhanced_image",)
    FUNCTION = "enhance"
    CATEGORY = "Oshun/Image Processing"

    def enhance(
        self,
        image: torch.Tensor,
        contrast: float,
        brightness: float,
        saturation: float,
        sharpness: float,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor]:
        """
        Apply image enhancements.

        Args:
            image: [B, H, W, C] tensor, values in [0, 1]
            contrast: Contrast multiplier
            brightness: Brightness offset
            saturation: Saturation multiplier
            sharpness: Sharpness amount
            mask: Optional mask for selective enhancement

        Returns:
            Enhanced image tensor
        """
        # Store original for masking
        original = image.clone()

        # Apply contrast
        mean = image.mean(dim=(1, 2), keepdim=True)
        image = (image - mean) * contrast + mean

        # Apply brightness
        image = image + brightness

        # Apply saturation
        gray = image.mean(dim=-1, keepdim=True)
        image = gray + (image - gray) * saturation

        # Apply sharpness
        if sharpness > 0:
            image = self._apply_sharpness(image, sharpness)

        # Clamp to valid range
        image = torch.clamp(image, 0, 1)

        # Apply mask if provided
        if mask is not None:
            mask = mask.unsqueeze(-1)  # Add channel dimension
            image = original * (1 - mask) + image * mask

        return (image,)

    def _apply_sharpness(
        self,
        image: torch.Tensor,
        amount: float
    ) -> torch.Tensor:
        """Apply unsharp mask for sharpening."""
        # Convert to BCHW for convolution
        image = image.permute(0, 3, 1, 2)

        # Gaussian blur kernel
        kernel_size = 3
        sigma = 1.0
        kernel = self._gaussian_kernel(kernel_size, sigma, image.device)
        kernel = kernel.expand(image.shape[1], 1, -1, -1)

        # Apply blur
        padding = kernel_size // 2
        blurred = F.conv2d(
            image,
            kernel,
            padding=padding,
            groups=image.shape[1]
        )

        # Unsharp mask
        sharpened = image + amount * (image - blurred)

        # Convert back to BHWC
        return sharpened.permute(0, 2, 3, 1)

    def _gaussian_kernel(
        self,
        size: int,
        sigma: float,
        device: torch.device
    ) -> torch.Tensor:
        """Create Gaussian kernel."""
        coords = torch.arange(size, device=device).float() - size // 2
        g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
        g = g / g.sum()
        return torch.outer(g, g).unsqueeze(0).unsqueeze(0)

class OshunPromptEnhancer:
    """
    Enhance prompts with quality boosters and style keywords.
    """

    QUALITY_BOOSTERS = [
        "highly detailed",
        "4k",
        "8k",
        "masterpiece",
        "best quality",
        "sharp focus",
        "intricate details",
    ]

    STYLE_PRESETS = {
        "photorealistic": [
            "photorealistic",
            "hyperrealistic",
            "photography",
            "DSLR",
            "raw photo",
        ],
        "artistic": [
            "artistic",
            "creative",
            "expressive",
            "painterly",
        ],
        "anime": [
            "anime style",
            "manga",
            "2D illustration",
            "cel shading",
        ],
        "cinematic": [
            "cinematic",
            "movie still",
            "dramatic lighting",
            "film grain",
        ],
    }

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": ""
                }),
                "style": (list(cls.STYLE_PRESETS.keys()),),
                "quality_boost": ("BOOLEAN", {"default": True}),
                "detail_level": (["low", "medium", "high", "ultra"],),
            },
            "optional": {
                "custom_keywords": ("STRING", {
                    "multiline": False,
                    "default": ""
                }),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("enhanced_prompt",)
    FUNCTION = "enhance_prompt"
    CATEGORY = "Oshun/Text Processing"

    def enhance_prompt(
        self,
        prompt: str,
        style: str,
        quality_boost: bool,
        detail_level: str,
        custom_keywords: str = ""
    ) -> Tuple[str]:
        """Enhance the prompt with additional keywords."""
        parts = [prompt.strip()]

        # Add style keywords
        if style in self.STYLE_PRESETS:
            parts.extend(self.STYLE_PRESETS[style][:2])

        # Add quality boosters
        if quality_boost:
            detail_map = {
                "low": 1,
                "medium": 2,
                "high": 4,
                "ultra": 6,
            }
            count = detail_map.get(detail_level, 2)
            parts.extend(self.QUALITY_BOOSTERS[:count])

        # Add custom keywords
        if custom_keywords:
            parts.extend([k.strip() for k in custom_keywords.split(",")])

        enhanced = ", ".join(parts)
        return (enhanced,)

import boto3
import requests
from typing import Tuple, Dict, Any
import numpy as np
from PIL import Image
import io
import os

class OshunS3Upload:
    """
    Upload generated images to S3.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "bucket": ("STRING", {"default": "oshun-outputs"}),
                "prefix": ("STRING", {"default": "generated/"}),
                "format": (["png", "jpg", "webp"],),
            },
            "optional": {
                "quality": ("INT", {
                    "default": 95,
                    "min": 1,
                    "max": 100,
                    "step": 1
                }),
                "metadata": ("STRING", {
                    "multiline": True,
                    "default": "{}"
                }),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("s3_urls", "s3_keys")
    FUNCTION = "upload"
    CATEGORY = "Oshun/Integration"

    def __init__(self):
        self.s3_client = boto3.client(
            "s3",
            aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
            aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
            region_name=os.environ.get("AWS_REGION", "us-east-1"),
        )

    def upload(
        self,
        images: torch.Tensor,
        bucket: str,
        prefix: str,
        format: str,
        quality: int = 95,
        metadata: str = "{}"
    ) -> Tuple[str, str]:
        """
        Upload images to S3.

        Returns:
            Tuple of (comma-separated URLs, comma-separated keys)
        """
        import json
        from datetime import datetime
        import uuid

        urls = []
        keys = []

        # Parse metadata
        try:
            meta_dict = json.loads(metadata)
        except json.JSONDecodeError:
            meta_dict = {}

        # Process each image in batch
        for i, img_tensor in enumerate(images):
            # Convert to PIL Image
            img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
            pil_img = Image.fromarray(img_np)

            # Encode to bytes
            buffer = io.BytesIO()
            save_format = "JPEG" if format == "jpg" else format.upper()
            save_kwargs = {"quality": quality} if format in ["jpg", "webp"] else {}
            pil_img.save(buffer, format=save_format, **save_kwargs)
            buffer.seek(0)

            # Generate key
            timestamp = datetime.utcnow().strftime("%Y/%m/%d")
            filename = f"{uuid.uuid4().hex}.{format}"
            key = f"{prefix.rstrip('/')}/{timestamp}/{filename}"

            # Upload
            content_type = {
                "png": "image/png",
                "jpg": "image/jpeg",
                "webp": "image/webp",
            }.get(format, "image/png")

            self.s3_client.upload_fileobj(
                buffer,
                bucket,
                key,
                ExtraArgs={
                    "ContentType": content_type,
                    "Metadata": {k: str(v) for k, v in meta_dict.items()},
                }
            )

            # Generate URL
            url = f"https://{bucket}.s3.amazonaws.com/{key}"
            urls.append(url)
            keys.append(key)

        return (",".join(urls), ",".join(keys))


class OshunWebhookNotify:
    """
    Send webhook notification when processing completes.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "trigger": ("*",),  # Any input triggers the webhook
                "webhook_url": ("STRING", {"default": ""}),
                "event_type": ("STRING", {"default": "generation_complete"}),
            },
            "optional": {
                "payload": ("STRING", {
                    "multiline": True,
                    "default": "{}"
                }),
                "include_image_urls": ("BOOLEAN", {"default": True}),
                "image_urls": ("STRING", {"default": ""}),
            },
        }

    RETURN_TYPES = ("BOOLEAN",)
    RETURN_NAMES = ("success",)
    FUNCTION = "notify"
    CATEGORY = "Oshun/Integration"
    OUTPUT_NODE = True  # This is an output node

    def notify(
        self,
        trigger: Any,
        webhook_url: str,
        event_type: str,
        payload: str = "{}",
        include_image_urls: bool = True,
        image_urls: str = ""
    ) -> Tuple[bool]:
        """Send webhook notification."""
        import json

        if not webhook_url:
            return (False,)

        try:
            # Build payload
            data = json.loads(payload)
            data["event_type"] = event_type
            data["timestamp"] = datetime.utcnow().isoformat()

            if include_image_urls and image_urls:
                data["image_urls"] = [u.strip() for u in image_urls.split(",")]

            # Send webhook
            response = requests.post(
                webhook_url,
                json=data,
                headers={"Content-Type": "application/json"},
                timeout=30,
            )
            response.raise_for_status()

            return (True,)

        except Exception as e:
            print(f"Webhook failed: {e}")
            return (False,)

@classmethod
def INPUT_TYPES(cls):
    return {
        "required": {
            # Text input
            "text": ("STRING", {
                "multiline": True,  # Multi-line text area
                "default": "",
                "placeholder": "Enter text...",
            }),

            # Integer with slider
            "steps": ("INT", {
                "default": 30,
                "min": 1,
                "max": 150,
                "step": 1,
                "display": "slider",  # or "number"
            }),

            # Float with slider
            "cfg": ("FLOAT", {
                "default": 7.5,
                "min": 1.0,
                "max": 30.0,
                "step": 0.5,
                "round": 2,  # Decimal places
                "display": "slider",
            }),

            # Dropdown selection
            "sampler": (["euler", "euler_ancestral", "dpmpp_2m"],),

            # Boolean toggle
            "enabled": ("BOOLEAN", {"default": True}),
        },
        "optional": {
            # Optional inputs (can be None)
            "mask": ("MASK",),
        },
        "hidden": {
            # Hidden inputs (not shown in UI)
            "node_id": "UNIQUE_ID",
            "prompt": "PROMPT",
        },
    }

# Define custom type
class OshunStyleData:
    """Custom data type for style information."""

    def __init__(self, style_name: str, weights: dict):
        self.style_name = style_name
        self.weights = weights


# Node that outputs custom type
class OshunStyleLoader:
    RETURN_TYPES = ("OSHUN_STYLE",)  # Custom type name
    RETURN_NAMES = ("style",)

    def load_style(self, name: str) -> Tuple[OshunStyleData]:
        style = OshunStyleData(name, {"contrast": 1.2})
        return (style,)


# Node that accepts custom type
class OshunApplyStyle:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "style": ("OSHUN_STYLE",),  # Accepts custom type
            },
        }

class OshunConditionalProcess:
    """Process only if condition is met."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "condition": ("BOOLEAN",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"

    # Enable lazy evaluation
    @classmethod
    def IS_CHANGED(cls, image, condition):
        """Return value that changes when re-execution needed."""
        # Return hash of inputs or None to always execute
        return hash((id(image), condition))

    def process(self, image, condition):
        if not condition:
            return (image,)  # Pass through unchanged

        # Process image
        return (self._apply_processing(image),)

from comfy.utils import ProgressBar

class OshunBatchProcess:
    """Process multiple items with progress."""

    FUNCTION = "process"

    def process(self, images, **kwargs):
        results = []
        pbar = ProgressBar(len(images))

        for i, image in enumerate(images):
            result = self._process_single(image, **kwargs)
            results.append(result)
            pbar.update_absolute(i + 1)

        return (torch.stack(results),)

import hashlib
from functools import lru_cache

class OshunCachedLoader:
    """Load resources with caching."""

    _cache = {}

    @classmethod
    def IS_CHANGED(cls, resource_name):
        # Only reload if file changed
        import os
        path = cls._get_path(resource_name)
        if os.path.exists(path):
            return os.path.getmtime(path)
        return None

    @classmethod
    @lru_cache(maxsize=10)
    def _load_resource(cls, resource_name: str):
        """Cached resource loading."""
        # Load and return resource
        pass

    def load(self, resource_name: str):
        return (self._load_resource(resource_name),)

class OshunDynamicInputs:
    """Node with dynamic number of inputs."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "count": ("INT", {"default": 2, "min": 1, "max": 10}),
            },
        }

    # Dynamic inputs added at runtime
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """Validate dynamic inputs."""
        count = kwargs.get("count", 2)
        for i in range(count):
            if f"image_{i}" not in kwargs:
                return f"Missing image_{i}"
        return True

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "combine"

    def combine(self, count, **kwargs):
        images = [kwargs[f"image_{i}"] for i in range(count)]
        return (torch.cat(images, dim=0),)

import pytest
import torch
from oshun_nodes.nodes.image_processing import OshunImageEnhance

class TestOshunImageEnhance:
    """Tests for OshunImageEnhance node."""

    @pytest.fixture
    def node(self):
        return OshunImageEnhance()

    @pytest.fixture
    def sample_image(self):
        # Create test image [B, H, W, C]
        return torch.rand(1, 64, 64, 3)

    def test_input_types(self):
        """Test INPUT_TYPES class method."""
        types = OshunImageEnhance.INPUT_TYPES()

        assert "required" in types
        assert "image" in types["required"]
        assert "contrast" in types["required"]

    def test_enhance_no_change(self, node, sample_image):
        """Test with neutral settings."""
        result = node.enhance(
            image=sample_image,
            contrast=1.0,
            brightness=0.0,
            saturation=1.0,
            sharpness=0.0,
        )

        assert len(result) == 1
        assert result[0].shape == sample_image.shape
        torch.testing.assert_close(result[0], sample_image, atol=1e-5, rtol=1e-5)

    def test_enhance_contrast(self, node, sample_image):
        """Test contrast adjustment."""
        result = node.enhance(
            image=sample_image,
            contrast=2.0,
            brightness=0.0,
            saturation=1.0,
            sharpness=0.0,
        )

        # Contrast should increase variance
        original_var = sample_image.var()
        result_var = result[0].var()
        assert result_var > original_var

    def test_enhance_with_mask(self, node, sample_image):
        """Test masked enhancement."""
        mask = torch.zeros(1, 64, 64)
        mask[:, 32:, :] = 1.0  # Mask bottom half

        result = node.enhance(
            image=sample_image,
            contrast=2.0,
            brightness=0.0,
            saturation=1.0,
            sharpness=0.0,
            mask=mask,
        )

        # Top half should be unchanged
        torch.testing.assert_close(
            result[0][:, :32, :, :],
            sample_image[:, :32, :, :],
            atol=1e-5,
            rtol=1e-5
        )

    def test_output_range(self, node, sample_image):
        """Test output is in valid range."""
        result = node.enhance(
            image=sample_image,
            contrast=3.0,
            brightness=0.5,
            saturation=2.0,
            sharpness=1.0,
        )

        assert result[0].min() >= 0.0
        assert result[0].max() <= 1.0

import pytest
from comfy.workflow import Workflow

class TestOshunNodesIntegration:
    """Integration tests with ComfyUI."""

    @pytest.fixture
    def comfy_api(self):
        """Connect to local ComfyUI."""
        from comfyui_api import ComfyUIAPI
        return ComfyUIAPI(host="localhost", port=8188)

    def test_workflow_execution(self, comfy_api):
        """Test node in a complete workflow."""
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {"image": "test.png"}
            },
            "2": {
                "class_type": "OshunImageEnhance",
                "inputs": {
                    "image": ["1", 0],
                    "contrast": 1.2,
                    "brightness": 0.1,
                    "saturation": 1.1,
                    "sharpness": 0.3,
                }
            },
            "3": {
                "class_type": "PreviewImage",
                "inputs": {"images": ["2", 0]}
            }
        }

        result = comfy_api.execute(workflow)
        assert result["status"] == "success"

torch>=2.0.0
numpy>=1.24.0
Pillow>=9.0.0
boto3>=1.26.0  # For S3 integration
requests>=2.28.0  # For webhooks

"""Installation script for Oshun custom nodes."""

import subprocess
import sys
import os

def install_requirements():
    """Install Python requirements."""
    requirements_path = os.path.join(
        os.path.dirname(__file__),
        "requirements.txt"
    )

    if os.path.exists(requirements_path):
        subprocess.check_call([
            sys.executable,
            "-m",
            "pip",
            "install",
            "-r",
            requirements_path,
        ])

def main():
    print("Installing Oshun Custom Nodes...")
    install_requirements()
    print("Installation complete!")

if __name__ == "__main__":
    main()

[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"

[project]
name = "oshun-comfyui-nodes"
version = "1.0.0"
description = "Oshun custom nodes for ComfyUI"
requires-python = ">=3.10"
dependencies = [
    "torch>=2.0.0",
    "numpy>=1.24.0",
    "Pillow>=9.0.0",
]

[project.optional-dependencies]
dev = [
    "pytest>=7.0.0",
    "pytest-cov>=4.0.0",
]
integration = [
    "boto3>=1.26.0",
    "requests>=2.28.0",
]

# docker/runpod/comfyui/Dockerfile
FROM comfyanonymous/comfyui:latest

# Copy custom nodes
COPY custom_nodes/oshun_nodes /app/ComfyUI/custom_nodes/oshun_nodes

# Install dependencies
RUN pip install -r /app/ComfyUI/custom_nodes/oshun_nodes/requirements.txt

# Verify installation
RUN python -c "from oshun_nodes import NODE_CLASS_MAPPINGS; print(f'Loaded {len(NODE_CLASS_MAPPINGS)} nodes')"

# .github/workflows/deploy-custom-nodes.yml
name: Deploy Custom Nodes

on:
  push:
    paths:
      - 'custom_nodes/oshun_nodes/**'

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v4
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: |
          pip install -r custom_nodes/oshun_nodes/requirements.txt
          pip install pytest pytest-cov

      - name: Run tests
        run: |
          pytest custom_nodes/oshun_nodes/tests/ -v --cov

  build:
    needs: test
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Build Docker image
        run: |
          docker build -t oshunai/runpod-comfyui:latest \
            -f docker/runpod/comfyui/Dockerfile .

      - name: Push to registry
        run: |
          docker push oshunai/runpod-comfyui:latest

# Test node availability
async def verify_custom_nodes():
    """Verify custom nodes are available on provider."""
    from comfyui_api import ComfyUIAPI

    api = ComfyUIAPI(endpoint_url=os.environ["COMFYUI_ENDPOINT"])

    # Get available nodes
    nodes = await api.get_node_info()

    # Check for Oshun nodes
    oshun_nodes = [n for n in nodes if n.startswith("Oshun")]

    expected_nodes = [
        "OshunImageEnhance",
        "OshunColorCorrect",
        "OshunPromptEnhancer",
        "OshunS3Upload",
    ]

    for node in expected_nodes:
        assert node in oshun_nodes, f"Missing node: {node}"

    print(f"All {len(expected_nodes)} custom nodes verified!")

from typing import Tuple, Optional
import torch

def process(
    self,
    image: torch.Tensor,
    strength: float,
    mask: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    """
    Process image with documented parameters.

    Args:
        image: Input tensor [B, H, W, C] in range [0, 1]
        strength: Effect strength in range [0, 2]
        mask: Optional mask tensor [B, H, W]

    Returns:
        Processed image tensor
    """
    pass

def process(self, image: torch.Tensor, **kwargs) -> Tuple[torch.Tensor]:
    # Validate input
    if image.dim() != 4:
        raise ValueError(f"Expected 4D tensor, got {image.dim()}D")

    if image.shape[-1] not in [1, 3, 4]:
        raise ValueError(f"Expected 1/3/4 channels, got {image.shape[-1]}")

    # Handle empty batch
    if image.shape[0] == 0:
        return (image,)

    # Handle single-channel (grayscale)
    if image.shape[-1] == 1:
        image = image.repeat(1, 1, 1, 3)

    return self._process_impl(image, **kwargs)

def process(self, image: torch.Tensor, **kwargs) -> Tuple[torch.Tensor]:
    # Process in chunks for large batches
    if image.shape[0] > 4:
        results = []
        for i in range(0, image.shape[0], 4):
            chunk = image[i:i+4]
            results.append(self._process_chunk(chunk, **kwargs))
            torch.cuda.empty_cache()  # Free memory
        return (torch.cat(results),)

    return self._process_chunk(image, **kwargs)

def process(self, image: torch.Tensor, model_name: str, **kwargs):
    try:
        model = self.load_model(model_name)
    except FileNotFoundError:
        raise ValueError(
            f"Model '{model_name}' not found. "
            f"Available models: {', '.join(self.list_models())}"
        )
    except Exception as e:
        raise RuntimeError(f"Failed to load model '{model_name}': {e}")

    return self._process_with_model(image, model, **kwargs)

class OshunImageEnhance:
    """
    Enhance image quality with multiple adjustable parameters.

    This node provides professional-grade image enhancement with
    control over contrast, brightness, saturation, and sharpness.

    Features:
        - Non-destructive processing
        - Optional mask support for selective enhancement
        - GPU-accelerated operations

    Example workflow:
        LoadImage → OshunImageEnhance → SaveImage

    Notes:
        - All adjustments preserve the [0, 1] range
        - Use mask input for localized adjustments
        - Sharpening uses unsharp mask technique
    """