diff --git a/__init__.py b/__init__.py index f6b6485..6c63b7f 100644 --- a/__init__.py +++ b/__init__.py @@ -2,7 +2,8 @@ import os import sys -from .nodes.tools import HighPassFilter +from .nodes.image_feeder import ImageFeeder +from .nodes.image_filters import HighPassFilter from .nodes.face_wrapper import FaceWrapper from .nodes.face_fit_and_restore import FaceFitAndRestore @@ -15,13 +16,15 @@ NODE_PATH = os.path.dirname(os.path.realpath(__file__)) NODE_CLASS_MAPPINGS = { "FaceFitAndRestore": FaceFitAndRestore, "FaceWrapper": FaceWrapper, - "HighPassFilter": HighPassFilter + "HighPassFilter": HighPassFilter, + "ImageFeeder": ImageFeeder } NODE_DISPLAY_NAME_MAPPINGS = { "FaceFitAndRestore": "Face Fit or Restore", "FaceWrapper": "Face Wrapper", - "HighPassFilter": "High Pass Filter (HPF)" + "HighPassFilter": "High Pass Filter (HPF)", + "ImageFeeder": "Image Feeder" } -__version__ = "1.0.1" +__version__ = "1.1.0" diff --git a/nodes/image_feeder.py b/nodes/image_feeder.py new file mode 100644 index 0000000..dddfcbd --- /dev/null +++ b/nodes/image_feeder.py @@ -0,0 +1,157 @@ +import os +import torch +import numpy as np +from PIL import Image +from typing import List, Dict, Union, Tuple + + +class ImageFeeder: + """ComfyUI node for feeding images from a directory.""" + + def __init__(self): + self.current_dir = None + self.image_files = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING", { + "default": "", + "multiline": False + }), + "frame_number": ("INT", { + "default": 0, + "min": 0, + "step": 1 + }), + } + } + + RETURN_TYPES = ("IMAGE", "DICT") + RETURN_NAMES = ("image", "frames_data") + FUNCTION = "feed_images" + CATEGORY = "image" + + def _load_image(self, image_path: str) -> torch.Tensor: + """ + Load an image from path and convert it to ComfyUI format + + Args: + image_path: Path to the image file + + Returns: + torch.Tensor: Image in ComfyUI format (B,H,W,C) normalized to 0-1 range + """ + try: + image = Image.open(image_path) + + # Convert to RGB if necessary + if image.mode != 'RGB': + image = image.convert('RGB') + + # Convert to numpy array and normalize + image_np = np.array(image).astype(np.float32) / 255.0 + + # Convert to torch tensor and add batch dimension + return torch.from_numpy(image_np).unsqueeze(0) + + except Exception as e: + print(f"Error loading image {image_path}: {str(e)}") + return None + + def _scan_directory(self, directory: str) -> List[str]: + """ + Scan directory for image files + + Args: + directory: Path to directory to scan + + Returns: + List of image file paths + """ + # List of supported image extensions + valid_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.webp'} + + image_files = [] + + try: + # Get absolute path + abs_dir = os.path.abspath(directory) + + if not os.path.exists(abs_dir): + print(f"Directory not found: {abs_dir}") + return image_files + + # Scan directory for image files + print(f"Scanning directory: {abs_dir}") + + for file in os.listdir(abs_dir): + # Check file extension + if os.path.splitext(file)[1].lower() in valid_extensions: + full_path = os.path.join(abs_dir, file) + image_files.append(full_path) + + # Sort files for consistent ordering + image_files.sort() + + print(f"Found {len(image_files)} image files") + + except Exception as e: + print(f"Error scanning directory {directory}: {str(e)}") + + return image_files + + def feed_images(self, directory: str, frame_number: int) -> Tuple[torch.Tensor, Dict]: + """ + Main processing function + + Args: + directory: Path to images directory + frame_number: Index of frame to return + + Returns: + Tuple containing: + - Selected frame as tensor + - Dictionary with frames data + """ + # Check if directory changed + if directory != self.current_dir: + print(f"New directory detected, scanning: {directory}") + self.image_files = self._scan_directory(directory) + self.current_dir = directory + + # Prepare return data with frame to file mapping + frames_data = { + "frames": { + idx: file_path + for idx, file_path in enumerate(self.image_files) + }, + "current_frame": frame_number + } + + # Handle empty directory case + if not self.image_files: + print("No images found in directory") + # Return empty image and data + empty_image = torch.zeros((1, 64, 64, 3)) + return (empty_image, frames_data) + + # Validate frame number + max_frame = len(self.image_files) - 1 + if frame_number > max_frame: + print(f"Warning: Requested frame {frame_number} is out of range. Using last available frame ({max_frame})") + frame_number = max_frame + elif frame_number < 0: + print(f"Warning: Requested frame {frame_number} is negative. Using first frame (0)") + frame_number = 0 + + # Load selected frame + selected_frame = self._load_image(self.image_files[frame_number]) + + if selected_frame is None: + print(f"Failed to load frame {frame_number}") + empty_image = torch.zeros((1, 64, 64, 3)) + return (empty_image, frames_data) + + return (selected_frame, frames_data) diff --git a/nodes/tools.py b/nodes/image_filters.py similarity index 88% rename from nodes/tools.py rename to nodes/image_filters.py index 3c0c7f8..c22dd7e 100644 --- a/nodes/tools.py +++ b/nodes/image_filters.py @@ -27,13 +27,13 @@ class HighPassFilter: RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) - FUNCTION = "apply_effect" + FUNCTION = "apply_hpf" CATEGORY = "Face Processor/Tools" - def apply_effect(self, image: torch.Tensor, blur_radius: int, blur_iterations: int, - blend_opacity: float, input_black: int, input_white: int, - gamma: float, show_histogram: bool) -> tuple[Tensor]: - """Main processing function implementing the AE-style effect workflow""" + def apply_hpf(self, image: torch.Tensor, blur_radius: int, blur_iterations: int, + blend_opacity: float, input_black: int, input_white: int, + gamma: float, show_histogram: bool) -> tuple[Tensor]: + """Main processing function implementing the High Pass Filter""" # Convert tensor to numpy array np_img = image[0].numpy() * 255 np_img = np.clip(np_img, 0, 255).astype(np.uint8)