custom_nodes/__init__.py

"""
DGG NWN Custom Nodes for ComfyUI
Specialized nodes for Neverwinter Nights character enhancement workflow.
"""

class NWNCharacterEnhancePrompt:
    """
    Generate optimized prompts for NWN character enhancement.
    Automatically creates prompt/negative prompt pairs based on character archetype.
    """
    
    CHARACTER_PRESETS = {
        "paladin_female": {
            "base": "female paladin warrior, ornate silver plate armor with golden trim",
            "features": "noble elven features, determined expression, flowing hair",
        },
        "paladin_male": {
            "base": "male paladin warrior, heavy plate armor with holy symbols",
            "features": "strong jawline, short hair, battle-worn face",
        },
        "mage_female": {
            "base": "female arcane mage, flowing magical robes with arcane patterns",
            "features": "intelligent eyes, mysterious aura, elegant features",
        },
        "mage_male": {
            "base": "male wizard, elaborate robes with mystical symbols",
            "features": "wise expression, aged features, long beard",
        },
        "rogue_female": {
            "base": "female rogue assassin, dark leather armor with hidden blades",
            "features": "cunning eyes, agile build, sharp features",
        },
        "rogue_male": {
            "base": "male thief, practical leather gear with pouches",
            "features": "shadowy expression, quick reflexes, scarred face",
        },
        "warrior_female": {
            "base": "female barbarian warrior, fur-lined armor with battle damage",
            "features": "fierce expression, muscular build, war paint",
        },
        "warrior_male": {
            "base": "male fighter, practical plate mail with sword and shield",
            "features": "battle-hardened face, strong build, short cropped hair",
        },
    }
    
    QUALITY_SUFFIXES = {
        "ultra": "masterpiece, best quality, ultra detailed, 8k uhd, ray tracing, photorealistic",
        "high": "high quality, detailed, 4k, professional lighting, realistic",
        "medium": "good quality, detailed, sharp focus",
    }
    
    NEGATIVE_BASE = "low quality, blurry, pixelated, low poly, cartoonish, anime style, deformed, bad anatomy, extra limbs, watermark, signature, text"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "character_class": (["paladin", "mage", "rogue", "warrior"],),
                "gender": (["female", "male"],),
                "race": (["human", "elf", "dwarf", "halfling", "half-elf", "half-orc"],),
                "quality": (["ultra", "high", "medium"],),
                "custom_details": ("STRING", {"default": "", "multiline": True}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("positive_prompt", "negative_prompt")
    FUNCTION = "generate_prompt"
    CATEGORY = "DGG/NWN"
    
    def generate_prompt(self, character_class, gender, race, quality, custom_details):
        # Get preset
        preset_key = f"{character_class}_{gender}"
        preset = self.CHARACTER_PRESETS.get(preset_key, self.CHARACTER_PRESETS["warrior_male"])
        
        # Build prompt
        race_desc = f"{race} " if race != "human" else ""
        
        positive = f"photorealistic {race_desc}{preset['base']}, {preset['features']}"
        
        if custom_details.strip():
            positive += f", {custom_details.strip()}"
        
        positive += f", {self.QUALITY_SUFFIXES[quality]}"
        
        negative = self.NEGATIVE_BASE
        
        return (positive, negative)


class NWNImagePreprocess:
    """
    Preprocess NWN screenshots for better enhancement results.
    - Removes UI elements
    - Normalizes lighting
    - Isolates character
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "remove_background": ("BOOLEAN", {"default": True}),
                "normalize_lighting": ("BOOLEAN", {"default": True}),
                "upscale_first": ("BOOLEAN", {"default": False}),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "preprocess"
    CATEGORY = "DGG/NWN"
    
    def preprocess(self, image, remove_background, normalize_lighting, upscale_first):
        import torch
        import numpy as np
        
        # Convert to numpy
        if isinstance(image, torch.Tensor):
            img_np = image.cpu().numpy()
            if img_np.ndim == 4:
                img_np = img_np[0]  # Remove batch dimension
        else:
            img_np = np.array(image)
        
        # Normalize to 0-1 if needed
        if img_np.max() > 1.0:
            img_np = img_np / 255.0
        
        if normalize_lighting:
            # Simple contrast normalization
            for c in range(3):
                channel = img_np[:, :, c]
                min_val = channel.min()
                max_val = channel.max()
                if max_val > min_val:
                    img_np[:, :, c] = (channel - min_val) / (max_val - min_val)
        
        # Convert back to tensor
        result = torch.from_numpy(img_np).unsqueeze(0)
        
        return (result,)


class NWNBatchProcessor:
    """
    Process multiple NWN character images in batch.
    Useful for generating front/side/back views consistently.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "prompt_template": ("STRING", {"default": "photorealistic fantasy character, {view} view"}),
                "views": ("STRING", {"default": "front,side,back"}),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "view_prompts")
    FUNCTION = "process_batch"
    CATEGORY = "DGG/NWN"
    
    def process_batch(self, images, prompt_template, views):
        view_list = [v.strip() for v in views.split(",")]
        prompts = [prompt_template.format(view=v) for v in view_list]
        return (images, "\n".join(prompts))


# Node mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "NWNCharacterEnhancePrompt": NWNCharacterEnhancePrompt,
    "NWNImagePreprocess": NWNImagePreprocess,
    "NWNBatchProcessor": NWNBatchProcessor,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "NWNCharacterEnhancePrompt": "🎮 NWN Character Prompt",
    "NWNImagePreprocess": "🎮 NWN Image Preprocess",
    "NWNBatchProcessor": "🎮 NWN Batch Processor",
}