Add files via upload

pixel_sort_node.py

@@ -0,0 +1,160 @@
+"""
+ComfyUI Pixel Sort Node
+Integrates ASDFPixelSort Processing script with ComfyUI workflows
+"""
+
+import subprocess
+import os
+import tempfile
+import torch
+import numpy as np
+from PIL import Image
+
+class PixelSortNode:
+    """
+    A ComfyUI node that applies pixel sorting effects using the ASDFPixelSort Processing script.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mode": (["white", "black", "bright", "dark"], {"default": "white"}),
+                "white_threshold": ("INT", {"default": -12345678, "min": -16777216, "max": 0, "step": 1000}),
+                "black_threshold": ("INT", {"default": -3456789, "min": -16777216, "max": 0, "step": 1000}),
+                "bright_threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
+                "dark_threshold": ("INT", {"default": 223, "min": 0, "max": 255, "step": 1}),
+            }
+        }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_pixel_sort"
+    CATEGORY = "image/effects"
+    
+    def __init__(self):
+        # Path to Processing installation
+        self.processing_path = "/root/Processing/processing-4.3/processing-java"
+        self.sketch_path = "/root/Processing/ASDFPixelSort"
+        
+    def apply_pixel_sort(self, image, mode, white_threshold, black_threshold, bright_threshold, dark_threshold):
+        """
+        Apply pixel sorting to the input image using the Processing script.
+        """
+        
+        # Convert mode string to number
+        mode_map = {"white": 0, "black": 1, "bright": 2, "dark": 3}
+        mode_num = mode_map[mode]
+        
+        # Create temporary files
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as input_file:
+            input_path = input_file.name
+        
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as output_file:
+            output_path = output_file.name
+        
+        try:
+            # Convert tensor to PIL Image and save
+            img_tensor = image[0]  # Remove batch dimension
+            img_array = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
+            img_pil = Image.fromarray(img_array)
+            img_pil.save(input_path)
+            
+            # Prepare command arguments
+            cmd = [
+                self.processing_path,
+                f"--sketch={self.sketch_path}",
+                "--run",
+                input_path,
+                output_path,
+                str(mode_num),
+                str(white_threshold),
+                str(black_threshold),
+                str(bright_threshold),
+                str(dark_threshold)
+            ]
+            
+            # Set up environment for headless operation
+            env = os.environ.copy()
+            env["DISPLAY"] = ":99"
+            
+            # Run the Processing script
+            print(f"Running pixel sort with mode: {mode}")
+            result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=120)
+            
+            if result.returncode != 0:
+                print(f"Processing script error: {result.stderr}")
+                raise RuntimeError(f"Pixel sort failed: {result.stderr}")
+            
+            # Load the processed image
+            if not os.path.exists(output_path):
+                raise RuntimeError("Output image was not created")
+                
+            img_sorted = Image.open(output_path).convert("RGB")
+            img_array = np.array(img_sorted).astype(np.float32) / 255.0
+            img_tensor = torch.from_numpy(img_array).unsqueeze(0)  # Add batch dimension
+            
+            print(f"Pixel sort completed successfully")
+            return (img_tensor,)
+            
+        except subprocess.TimeoutExpired:
+            raise RuntimeError("Pixel sort operation timed out")
+        except Exception as e:
+            raise RuntimeError(f"Pixel sort failed: {str(e)}")
+        finally:
+            # Clean up temporary files
+            try:
+                if os.path.exists(input_path):
+                    os.unlink(input_path)
+                if os.path.exists(output_path):
+                    os.unlink(output_path)
+            except:
+                pass
+
+class PixelSortAdvancedNode:
+    """
+    Advanced pixel sort node with more control options.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mode": (["white", "black", "bright", "dark"], {"default": "white"}),
+            },
+            "optional": {
+                "white_threshold": ("INT", {"default": -12345678, "min": -16777216, "max": 0, "step": 1000}),
+                "black_threshold": ("INT", {"default": -3456789, "min": -16777216, "max": 0, "step": 1000}),
+                "bright_threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
+                "dark_threshold": ("INT", {"default": 223, "min": 0, "max": 255, "step": 1}),
+                "custom_sketch_path": ("STRING", {"default": ""}),
+            }
+        }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_pixel_sort_advanced"
+    CATEGORY = "image/effects"
+    
+    def apply_pixel_sort_advanced(self, image, mode, white_threshold=-12345678, black_threshold=-3456789, 
+                                bright_threshold=127, dark_threshold=223, custom_sketch_path=""):
+        """
+        Advanced pixel sorting with optional custom sketch path.
+        """
+        node = PixelSortNode()
+        if custom_sketch_path:
+            node.sketch_path = custom_sketch_path
+        
+        return node.apply_pixel_sort(image, mode, white_threshold, black_threshold, 
+                                   bright_threshold, dark_threshold)
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "PixelSortNode": PixelSortNode,
+    "PixelSortAdvancedNode": PixelSortAdvancedNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PixelSortNode": "Pixel Sort (ASDF)",
+    "PixelSortAdvancedNode": "Pixel Sort Advanced (ASDF)"
+}