diff --git a/README.md b/README.md index c0ffef7..ffee311 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,13 @@ - ✅ Нода аналитики тегов (top-k + JSON статистика) ## Ноды +- **4 WD14 v3 Models**: ViT, SwinV2, EVA02, ConvNeXT +- **Auto Download**: Models download automatically from Hugging Face +- **GPU Support**: CUDA acceleration for faster processing +- **Batch Processing**: Process multiple images from folders +- **Format Support**: JPG, JPEG, PNG, WEBP +- **Custom Tags**: Add/remove tags as needed +- **ComfyUI IMAGE Compatibility**: Uses native tensor image format for modern ComfyUI builds ### 1) BSS Load Images from Folder 📂 Загружает изображения из папки (JPG/JPEG/PNG/WEBP) и отдает `IMAGE` + имена файлов. @@ -40,6 +47,39 @@ ### 5) BSS Save Captions 💾 Сохраняет caption в `txt/json/csv`, есть `overwrite` и `suffix`. +1. Use **BSS Load Images from Folder** to load your images +2. Connect to **BSS WD14 Batch Tagger** for each image +3. Set output folder for tag files +4. Run the workflow + +### Settings + +- **Model**: Choose WD14 model (auto-downloads if needed) +- **Threshold**: Tag confidence (0.35 default) +- **Character Threshold**: Separate threshold for character tags (WD category 4) +- **GPU**: Enable for faster processing +- **Prepend/Exclude**: Add custom tags or remove unwanted ones + +## ComfyUI Compatibility + +- Tested against current ComfyUI custom node API style (`NODE_CLASS_MAPPINGS`, `INPUT_TYPES`, `RETURN_TYPES`). +- Uses ComfyUI-native `IMAGE` tensors (`float32`, range `0..1`) in the loader node output. +- Tagger node accepts both tensor images (ComfyUI-native) and numpy arrays for backward compatibility. + +## Models + +- **WD ViT Tagger v3**: Fast, good quality (default) +- **WD SwinV2 Tagger v3**: Balanced speed/quality +- **WD EVA02 Large Tagger v3**: Best accuracy +- **WD ConvNeXT Tagger v3**: Modern architecture + +Models download automatically on first use. + +## Requirements + +- Python 3.8+ +- ComfyUI +- CUDA GPU (optional) ### 6) BSS Tag Analytics 📊 Считает top-k частых тегов и возвращает JSON статистику. @@ -67,3 +107,9 @@ pip install -r requirements.txt ## License MIT + +MIT License - see LICENSE file for details. + +--- + +**Author**: Blacksnowskill diff --git a/nodes.py b/nodes.py index f737c34..66a381f 100644 --- a/nodes.py +++ b/nodes.py @@ -208,6 +208,9 @@ def load_model_once(model_name: str, use_gpu: bool) -> bool: logger.error(f"CSV file not found: {csv_path}") return False + + tags = [] + tag_categories = [] with open(csv_path, newline='', encoding='utf-8') as f: reader = csv.reader(f) next(reader) @@ -219,6 +222,8 @@ def load_model_once(model_name: str, use_gpu: bool) -> bool: except (ValueError, IndexError): tag_categories.append(-1) + + # Cache the model and tags model_cache['session'] = sess model_cache['tags'] = tags model_cache['tag_categories'] = tag_categories @@ -335,6 +340,15 @@ class BSS_LoadImagesFolder: return True def load_images(self, folder_path: str) -> Tuple[List[torch.Tensor], List[str], str]: + """ + Load images from specified folder path. + + Args: + folder_path: Path to folder containing images + + Returns: + Tuple of (images, filenames, folder_path) + """ if not folder_path or not os.path.exists(folder_path): logger.error(f"Invalid folder path: {folder_path}") return [], [], folder_path @@ -465,6 +479,108 @@ class BSS_WD14BatchTagger: use_gpu, ) + def tag_batch(self, image: Any, filename: str, folder_path: str, model: str, + threshold: float, character_threshold: float, replace_underscore: bool, + use_gpu: bool, prepend_tags: str, exclude_tags: str) -> Tuple[str]: + """ + Tag a single image using WD14 model and save tags to txt file. + + Args: + image: Input image as numpy array + filename: Original filename + folder_path: Output folder path + model: Model name (format: "model_name|display_name") + threshold: Confidence threshold for tags + character_threshold: Character-specific threshold + replace_underscore: Whether to replace underscores with spaces + use_gpu: Whether to use GPU acceleration + prepend_tags: Tags to prepend to result + exclude_tags: Tags to exclude from result + + Returns: + Tuple containing the generated tags string + """ + try: + # Extract model name from dropdown selection (handle both formats) + if "|" in model: + model_name = model.split("|")[0] + else: + model_name = model + + # Load model once (cached for subsequent calls) + if not load_model_once(model_name, use_gpu): + logger.error(f"Failed to load model: {model_name}") + return ("",) + + # Get cached model data + sess = model_cache['session'] + tags = model_cache['tags'] + tag_categories = model_cache.get('tag_categories', []) + input_name = model_cache['input_name'] + input_size = model_cache['input_size'] + + # Process tags (replace underscores if needed) + processed_tags = [] + for tag in tags: + processed_tag = tag.replace("_", " ") if replace_underscore else tag + processed_tags.append(processed_tag) + + # Parse exclude tags + exclude_list = [t.strip().lower() for t in exclude_tags.split(",") if t.strip()] + + # Process image + if isinstance(image, torch.Tensor): + image_np = image.detach().cpu().numpy() + else: + image_np = np.asarray(image) + + if image_np is None or image_np.size == 0: + logger.warning(f"Empty image array for {filename}, skipping") + return ("",) + + if image_np.ndim == 4: + image_np = image_np[0] + + if image_np.shape[-1] == 4: + image_np = image_np[:, :, :3] + + if image_np.dtype != np.uint8: + image_np = np.clip(image_np, 0.0, 1.0) + image_np = (image_np * 255.0).astype(np.uint8) + + img = Image.fromarray(image_np) + ratio = input_size / max(img.size) + new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio)) + + if new_size[0] <= 0 or new_size[1] <= 0: + logger.error(f"Invalid resize dimensions {new_size} for {filename}") + return ("",) + + # Resize and center image + try: + # Use new Pillow API if available + img = img.resize(new_size, Image.Resampling.LANCZOS) + except AttributeError: + # Fallback for older Pillow versions + img = img.resize(new_size, Image.LANCZOS) + square = Image.new("RGB", (input_size, input_size), (255, 255, 255)) + paste_pos = ((input_size - new_size[0]) // 2, (input_size - new_size[1]) // 2) + square.paste(img, paste_pos) + + # Convert to model input format (BGR, normalized) + inp = np.expand_dims(np.array(square).astype(np.float32)[:, :, ::-1], 0) + + # Run inference + probs = sess.run(None, {input_name: inp})[0][0] + + # Filter tags by threshold and exclusions + result_tags = [] + for idx, (tag, prob) in enumerate(zip(processed_tags, probs)): + current_threshold = character_threshold if idx < len(tag_categories) and tag_categories[idx] == 4 else threshold + if prob > current_threshold and tag.lower() not in exclude_list: + result_tags.append(tag) + + # Build output string output_tags = prepend_tags.strip() if output_tags and result_tags: output_tags += ", "