From 2bf83f8ae40b6d6800aef19dd23783044521580b Mon Sep 17 00:00:00 2001 From: kadirnar Date: Sun, 8 Feb 2026 15:54:29 +0000 Subject: [PATCH] Add new pipelines for multimodal and audio tasks, enhance NLP capabilities --- .../audio/audio_classification/node.py | 29 +++++ .../automatic_speech_recognition/node.py | 32 +++++ huggingface_tasks/audio/text2speech/node.py | 41 ++++++ .../zero_shot_audio_classification/node.py | 30 +++++ .../image_classification/node.py | 52 +++----- .../image_feature_extraction/node.py | 36 ++++++ .../image_segmentation/node.py | 60 ++++----- .../computer_vision/image_to_image/node.py | 39 ++++++ .../keypoint_detection/node.py | 46 +++++++ .../computer_vision/mask_generation/node.py | 48 +++++++ .../computer_vision/object_detection/node.py | 73 +++++------ .../video_classification/node.py | 29 +++++ .../zero_shot_image_classification/node.py | 35 ++++++ .../zero_shot_object_detection/node.py | 48 +++++++ .../document_question_answering/node.py | 34 +++++ .../multimodal/feature_extraction/node.py | 31 +++++ .../multimodal/image2text/node.py | 33 +++++ .../multimodal/image_text_to_text/node.py | 38 ++++++ .../visual_question_answering/node.py | 34 +++++ .../conversational/node.py | 40 ++++++ .../fill_mask/node.py | 29 +++++ .../question_answering/node.py | 28 +++++ .../sentence_similarity/node.py | 32 +++++ .../summarization/node.py | 29 +++++ .../table_question_answering/node.py | 30 +++++ .../text2text_generation/node.py | 28 +++++ .../text_classification/node.py | 28 +++++ .../text_generation/node.py | 29 +++++ .../token_classification/node.py | 37 ++++++ .../translation/node.py | 28 +++++ .../zero_shot_classification/node.py | 30 +++++ huggingface_tasks/utils.py | 36 ++++++ nodes.py | 119 ++++++++++++++++-- requirements.txt | 13 +- 34 files changed, 1174 insertions(+), 130 deletions(-) create mode 100644 huggingface_tasks/audio/zero_shot_audio_classification/node.py create mode 100644 huggingface_tasks/computer_vision/image_feature_extraction/node.py create mode 100644 huggingface_tasks/computer_vision/image_to_image/node.py create mode 100644 huggingface_tasks/computer_vision/keypoint_detection/node.py create mode 100644 huggingface_tasks/computer_vision/mask_generation/node.py create mode 100644 huggingface_tasks/computer_vision/video_classification/node.py create mode 100644 huggingface_tasks/computer_vision/zero_shot_image_classification/node.py create mode 100644 huggingface_tasks/computer_vision/zero_shot_object_detection/node.py create mode 100644 huggingface_tasks/multimodal/image_text_to_text/node.py create mode 100644 huggingface_tasks/natural_language_processing/text2text_generation/node.py create mode 100644 huggingface_tasks/utils.py diff --git a/huggingface_tasks/audio/audio_classification/node.py b/huggingface_tasks/audio/audio_classification/node.py index e69de29..b340f1f 100644 --- a/huggingface_tasks/audio/audio_classification/node.py +++ b/huggingface_tasks/audio/audio_classification/node.py @@ -0,0 +1,29 @@ +from transformers import pipeline as hf_pipeline +import json + +audio_classification_model_list = [ + "MIT/ast-finetuned-audioset-10-10-0.4593", + "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition", +] + + +class AudioClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_path": ("STRING", {"default": ""}), + "model_name": (audio_classification_model_list, {"default": audio_classification_model_list[0]}), + "top_k": ("INT", {"default": 5, "min": 1, "max": 20}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("results_json",) + FUNCTION = "run_audio_classification" + CATEGORY = "Transformers/Audio/AudioClassification" + + def run_audio_classification(self, audio_path, model_name, top_k): + pipe = hf_pipeline("audio-classification", model=model_name) + result = pipe(audio_path, top_k=top_k) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/audio/automatic_speech_recognition/node.py b/huggingface_tasks/audio/automatic_speech_recognition/node.py index e69de29..eae6c98 100644 --- a/huggingface_tasks/audio/automatic_speech_recognition/node.py +++ b/huggingface_tasks/audio/automatic_speech_recognition/node.py @@ -0,0 +1,32 @@ +from transformers import pipeline as hf_pipeline + +asr_model_list = [ + "openai/whisper-tiny", + "openai/whisper-base", + "openai/whisper-small", +] + + +class AutomaticSpeechRecognitionPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_path": ("STRING", {"default": ""}), + "model_name": (asr_model_list, {"default": asr_model_list[0]}), + "language": ("STRING", {"default": "english"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("transcribed_text",) + FUNCTION = "run_asr" + CATEGORY = "Transformers/Audio/SpeechRecognition" + + def run_asr(self, audio_path, model_name, language): + generate_kwargs = {} + if language: + generate_kwargs["language"] = language + pipe = hf_pipeline("automatic-speech-recognition", model=model_name) + result = pipe(audio_path, generate_kwargs=generate_kwargs) + return (result["text"],) diff --git a/huggingface_tasks/audio/text2speech/node.py b/huggingface_tasks/audio/text2speech/node.py index e69de29..be2bd02 100644 --- a/huggingface_tasks/audio/text2speech/node.py +++ b/huggingface_tasks/audio/text2speech/node.py @@ -0,0 +1,41 @@ +from transformers import pipeline as hf_pipeline +import scipy.io.wavfile +import numpy as np +import tempfile +import os + +text_to_speech_model_list = [ + "suno/bark-small", + "facebook/mms-tts-eng", +] + + +class TextToSpeechPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (text_to_speech_model_list, {"default": text_to_speech_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("audio_path",) + FUNCTION = "run_tts" + CATEGORY = "Transformers/Audio/TextToSpeech" + + def run_tts(self, text, model_name): + pipe = hf_pipeline("text-to-speech", model=model_name) + result = pipe(text) + + audio_data = result["audio"] + sampling_rate = result["sampling_rate"] + + if isinstance(audio_data, np.ndarray): + if audio_data.dtype == np.float32 or audio_data.dtype == np.float64: + audio_data = (audio_data * 32767).astype(np.int16) + + output_path = os.path.join(tempfile.gettempdir(), "comfyui_tts_output.wav") + scipy.io.wavfile.write(output_path, rate=sampling_rate, data=audio_data) + return (output_path,) diff --git a/huggingface_tasks/audio/zero_shot_audio_classification/node.py b/huggingface_tasks/audio/zero_shot_audio_classification/node.py new file mode 100644 index 0000000..de951b1 --- /dev/null +++ b/huggingface_tasks/audio/zero_shot_audio_classification/node.py @@ -0,0 +1,30 @@ +from transformers import pipeline as hf_pipeline +import json + +zero_shot_audio_classification_model_list = [ + "laion/clap-htsat-unfused", + "laion/larger_clap_general", +] + + +class ZeroShotAudioClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_path": ("STRING", {"default": ""}), + "candidate_labels": ("STRING", {"default": "speech, music, noise"}), + "model_name": (zero_shot_audio_classification_model_list, {"default": zero_shot_audio_classification_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("results_json",) + FUNCTION = "run_zero_shot_audio" + CATEGORY = "Transformers/Audio/ZeroShotAudioClassification" + + def run_zero_shot_audio(self, audio_path, candidate_labels, model_name): + labels = [l.strip() for l in candidate_labels.split(",")] + pipe = hf_pipeline("zero-shot-audio-classification", model=model_name) + result = pipe(audio_path, candidate_labels=labels) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/computer_vision/image_classification/node.py b/huggingface_tasks/computer_vision/image_classification/node.py index f7b99f8..d23bf76 100644 --- a/huggingface_tasks/computer_vision/image_classification/node.py +++ b/huggingface_tasks/computer_vision/image_classification/node.py @@ -1,15 +1,16 @@ -from transformers import AutoImageProcessor, ResNetForImageClassification -import torch +from transformers import pipeline as hf_pipeline +import numpy as np from PIL import Image -from typing import Tuple image_classification_model_name_list = [ - "microsoft/resnet-50" + "microsoft/resnet-50", + "google/vit-base-patch16-224", ] + class ImageClassificationPipeline: @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), @@ -17,36 +18,17 @@ class ImageClassificationPipeline: }, } - RETURN_TYPES = ("STRING","FLOAT",) - FUNCTION = "classify_pipeline" + RETURN_TYPES = ("STRING", "FLOAT",) + RETURN_NAMES = ("label", "score",) + FUNCTION = "classify_image" + CATEGORY = "Transformers/ComputerVision/ImageClassification" - CATEGORY = "ComputerVision/Transformers" + def classify_image(self, image, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) - def classify_image(image: Image.Image, model_name: str) -> Tuple[str, float]: - """ - Classifies a given image using the specified model and returns the name and score of the predicted label. + pipe = hf_pipeline("image-classification", model=model_name) + result = pipe(pil_image) - Args: - image (PIL.Image.Image): The image to be classified. - model_name (str): The name of the model to be used for classification. - - Returns: - Tuple[str, float]: The name of the predicted label and its corresponding score. - """ - - # Load the specified model and corresponding image processor - image_processor = AutoImageProcessor.from_pretrained(model_name) - model = ResNetForImageClassification.from_pretrained(model_name) - - # Process the image and prepare it for the model - inputs = image_processor(image, return_tensors="pt") - - with torch.no_grad(): - logits = model(**inputs).logits - - # Predict the class with the highest score - predicted_label_id = logits.argmax(-1).item() - predicted_label_name = model.config.id2label[predicted_label_id] - predicted_score = torch.softmax(logits, -1)[0, predicted_label_id].item() - - return (predicted_label_name, predicted_score,) + top = result[0] + return (top["label"], top["score"],) diff --git a/huggingface_tasks/computer_vision/image_feature_extraction/node.py b/huggingface_tasks/computer_vision/image_feature_extraction/node.py new file mode 100644 index 0000000..ea3e014 --- /dev/null +++ b/huggingface_tasks/computer_vision/image_feature_extraction/node.py @@ -0,0 +1,36 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np +import json + +image_feature_extraction_model_list = [ + "google/vit-base-patch16-224", + "facebook/dinov2-small", +] + + +class ImageFeatureExtractionPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": (image_feature_extraction_model_list, {"default": image_feature_extraction_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("features_json",) + FUNCTION = "run_feature_extraction" + CATEGORY = "Transformers/ComputerVision/ImageFeatureExtraction" + + def run_feature_extraction(self, image, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("image-feature-extraction", model=model_name) + result = pipe(pil_image) + + features = np.array(result[0]).tolist() + output = {"shape": list(np.array(result[0]).shape), "features_preview": features[:5] if isinstance(features[0], list) else features[:10]} + return (json.dumps(output, indent=2),) diff --git a/huggingface_tasks/computer_vision/image_segmentation/node.py b/huggingface_tasks/computer_vision/image_segmentation/node.py index d691ba0..5a5547c 100644 --- a/huggingface_tasks/computer_vision/image_segmentation/node.py +++ b/huggingface_tasks/computer_vision/image_segmentation/node.py @@ -1,57 +1,45 @@ -from transformers import pipeline +from transformers import pipeline as hf_pipeline from PIL import Image +import numpy as np +import torch image_segmentation_model_name_list = [ "mattmdjaga/segformer_b2_clothes", - "nvidia/segformer-b1-finetuned-cityscapes-1024-1024" + "nvidia/segformer-b1-finetuned-cityscapes-1024-1024", ] + class ImageSegmentationPipeline: @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "category_name": ("STRING",), + "category_name": ("STRING", {"default": ""}), "model_name": (image_segmentation_model_name_list, {"default": image_segmentation_model_name_list[0]}), }, } RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("mask",) FUNCTION = "image_segmentation_pipeline" + CATEGORY = "Transformers/ComputerVision/ImageSegmentation" - CATEGORY = "ComputerVision/Transformers" + def image_segmentation_pipeline(self, image, category_name, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + pipe = hf_pipeline(task="image-segmentation", model=model_name) + result = pipe(pil_image) -def image_segmentation_pipeline(image_path: str, category_name: str, model_name: str) -> Image.Image: - """ - Given an image, a category name, and a model name, extracts and returns the mask for the specified category. + for item in result: + if item["label"] == category_name: + mask = item["mask"] + arr = np.array(mask).astype(np.float32) / 255.0 + if len(arr.shape) == 2: + arr = np.stack([arr, arr, arr], axis=-1) + return (torch.from_numpy(arr)[None,],) - Args: - image_path (str): The path to the image file. - category_name (str): The name of the category for which the mask is to be extracted. - model_name (str): The name of the model to be used for image segmentation. - - Returns: - PIL.Image.Image: The mask of the specified category, if found. Otherwise, returns None. - - Example: - >>> image_path = "path/to/image.jpg" - >>> category_name = "car" - >>> model_name = "nvidia/segformer-b1-finetuned-cityscapes-1024-1024" - >>> mask = image_segmentation_pipeline(image_path, category_name, model_name) - """ - - # Create the pipeline with the specified model - pipe = pipeline(task="image-segmentation", model=model_name) - - # Process the image - result = pipe(image_path) - - # Search for the mask of the specified category - for item in result: - if item['label'] == category_name: - return (item['mask'],) - - # If the category is not found, return None - return (None,) + h, w = pil_image.size[1], pil_image.size[0] + empty = torch.zeros(1, h, w, 3) + return (empty,) diff --git a/huggingface_tasks/computer_vision/image_to_image/node.py b/huggingface_tasks/computer_vision/image_to_image/node.py new file mode 100644 index 0000000..feb9ff1 --- /dev/null +++ b/huggingface_tasks/computer_vision/image_to_image/node.py @@ -0,0 +1,39 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np +import torch + +image_to_image_model_list = [ + "caidas/swin2SR-classical-sr-x2-64", + "caidas/swin2SR-lightweight-x2-64", +] + + +class ImageToImagePipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": (image_to_image_model_list, {"default": image_to_image_model_list[0]}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "run_image_to_image" + CATEGORY = "Transformers/ComputerVision/ImageToImage" + + def run_image_to_image(self, image, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("image-to-image", model=model_name) + result = pipe(pil_image) + + arr = np.array(result).astype(np.float32) / 255.0 + if len(arr.shape) == 2: + arr = np.stack([arr, arr, arr], axis=-1) + elif arr.shape[-1] == 4: + arr = arr[:, :, :3] + return (torch.from_numpy(arr)[None,],) diff --git a/huggingface_tasks/computer_vision/keypoint_detection/node.py b/huggingface_tasks/computer_vision/keypoint_detection/node.py new file mode 100644 index 0000000..5791240 --- /dev/null +++ b/huggingface_tasks/computer_vision/keypoint_detection/node.py @@ -0,0 +1,46 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image, ImageDraw +import numpy as np +import torch +import json + +keypoint_detection_model_list = [ + "magic-leap-community/superpoint", +] + + +class KeypointDetectionPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": (keypoint_detection_model_list, {"default": keypoint_detection_model_list[0]}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("annotated_image", "keypoints_json",) + FUNCTION = "run_keypoint_detection" + CATEGORY = "Transformers/ComputerVision/KeypointDetection" + + def run_keypoint_detection(self, image, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("keypoint-detection", model=model_name) + result = pipe(pil_image) + + draw = ImageDraw.Draw(pil_image) + keypoints_data = [] + if result and len(result) > 0: + for detection in result: + keypoints = detection.get("keypoints", []) + for kp in keypoints: + x, y = kp[0], kp[1] + draw.ellipse((x - 2, y - 2, x + 2, y + 2), fill="red") + keypoints_data.append({"x": float(x), "y": float(y)}) + + arr = np.array(pil_image).astype(np.float32) / 255.0 + tensor_image = torch.from_numpy(arr)[None,] + return (tensor_image, json.dumps(keypoints_data, indent=2),) diff --git a/huggingface_tasks/computer_vision/mask_generation/node.py b/huggingface_tasks/computer_vision/mask_generation/node.py new file mode 100644 index 0000000..09ac05a --- /dev/null +++ b/huggingface_tasks/computer_vision/mask_generation/node.py @@ -0,0 +1,48 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np +import torch +import json + +mask_generation_model_list = [ + "facebook/sam-vit-base", + "facebook/sam-vit-large", +] + + +class MaskGenerationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": (mask_generation_model_list, {"default": mask_generation_model_list[0]}), + "points_per_batch": ("INT", {"default": 64, "min": 1, "max": 256}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("overlay_image", "masks_json",) + FUNCTION = "run_mask_generation" + CATEGORY = "Transformers/ComputerVision/MaskGeneration" + + def run_mask_generation(self, image, model_name, points_per_batch): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("mask-generation", model=model_name, points_per_batch=points_per_batch) + result = pipe(pil_image) + + overlay = np.array(pil_image).astype(np.float32) + masks_info = [] + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)] + for i, mask in enumerate(result["masks"][:10]): + mask_arr = np.array(mask) + color = colors[i % len(colors)] + for c in range(3): + overlay[:, :, c] = np.where(mask_arr, overlay[:, :, c] * 0.5 + color[c] * 0.5, overlay[:, :, c]) + masks_info.append({"mask_index": i, "score": float(result["scores"][i]) if "scores" in result else None}) + + overlay = np.clip(overlay, 0, 255).astype(np.float32) / 255.0 + tensor_image = torch.from_numpy(overlay)[None,] + return (tensor_image, json.dumps(masks_info, indent=2),) diff --git a/huggingface_tasks/computer_vision/object_detection/node.py b/huggingface_tasks/computer_vision/object_detection/node.py index e825d8b..8d39de9 100644 --- a/huggingface_tasks/computer_vision/object_detection/node.py +++ b/huggingface_tasks/computer_vision/object_detection/node.py @@ -1,65 +1,48 @@ -from transformers import AutoImageProcessor, AutoModelForObjectDetection +from transformers import pipeline as hf_pipeline from PIL import Image, ImageDraw +import numpy as np import torch - +import json object_detection_model_name_list = [ - "mattmdjaga/segformer_b2_clothes", - "nvidia/segformer-b1-finetuned-cityscapes-1024-1024" + "facebook/detr-resnet-50", + "hustvl/yolos-tiny", ] + class ObjectDetectionPipeline: @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "category_name": ("STRING",), "model_name": (object_detection_model_name_list, {"default": object_detection_model_name_list[0]}), - "threshold": ("FLOAT", {"default": 0.5}), + "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}), }, } - RETURN_TYPES = ("IMAGE",) - FUNCTION = "image_segmentation_pipeline" + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("annotated_image", "detections_json",) + FUNCTION = "object_detection_pipeline" + CATEGORY = "Transformers/ComputerVision/ObjectDetection" - CATEGORY = "ComputerVision/Transformers" + def object_detection_pipeline(self, image, model_name, threshold): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + pipe = hf_pipeline("object-detection", model=model_name) + results = pipe(pil_image, threshold=threshold) -def object_detection_pipeline(image_path, model_name="facebook/detr-resnet-50", threshold=0.5): - """ - Detects objects in an image using a pre-trained object detection model, draws bounding boxes around them, - and prints the labels and scores of detected objects. + draw = ImageDraw.Draw(pil_image) + for det in results: + box = det["box"] + x, y, x2, y2 = box["xmin"], box["ymin"], box["xmax"], box["ymax"] + draw.rectangle((x, y, x2, y2), outline="red", width=2) + label_text = f"{det['label']} ({det['score']:.2f})" + draw.text((x, y), label_text, fill="white") - Args: - image_path (str): The path to the input image. - model_name (str): The name of the pre-trained object detection model. - threshold (float): The confidence threshold for object detection. + arr = np.array(pil_image).astype(np.float32) / 255.0 + tensor_image = torch.from_numpy(arr)[None,] + detections_json = json.dumps(results, indent=2, default=str) - Returns: - PIL.Image.Image: The image with bounding boxes drawn around detected objects. - """ - - # Load the image - image = Image.open(image_path) - - # Load the image processor and model - image_processor = AutoImageProcessor.from_pretrained(model_name, revision="no_timm") - model = AutoModelForObjectDetection.from_pretrained(model_name, revision="no_timm") - - with torch.no_grad(): - inputs = image_processor(images=image, return_tensors="pt") - outputs = model(**inputs) - target_sizes = torch.tensor([image.size[::-1]]) - results = image_processor.post_process_object_detection(outputs, threshold=threshold, target_sizes=target_sizes)[0] - - # Draw bounding boxes, labels, and scores on the image - draw = ImageDraw.Draw(image) - for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): - box = [round(i, 2) for i in box.tolist()] - x, y, x2, y2 = tuple(box) - draw.rectangle((x, y, x2, y2), outline="red", width=1) - label_text = f"{model.config.id2label[label.item()]} ({score:.2f})" - draw.text((x, y), label_text, fill="white") - - return (image,) + return (tensor_image, detections_json,) diff --git a/huggingface_tasks/computer_vision/video_classification/node.py b/huggingface_tasks/computer_vision/video_classification/node.py new file mode 100644 index 0000000..a34939b --- /dev/null +++ b/huggingface_tasks/computer_vision/video_classification/node.py @@ -0,0 +1,29 @@ +from transformers import pipeline as hf_pipeline +import json + +video_classification_model_list = [ + "MCG-NJU/videomae-base-finetuned-kinetics", + "MCG-NJU/videomae-small-finetuned-kinetics", +] + + +class VideoClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "video_path": ("STRING", {"default": ""}), + "model_name": (video_classification_model_list, {"default": video_classification_model_list[0]}), + "top_k": ("INT", {"default": 5, "min": 1, "max": 20}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("results_json",) + FUNCTION = "run_video_classification" + CATEGORY = "Transformers/ComputerVision/VideoClassification" + + def run_video_classification(self, video_path, model_name, top_k): + pipe = hf_pipeline("video-classification", model=model_name) + result = pipe(video_path, top_k=top_k) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/computer_vision/zero_shot_image_classification/node.py b/huggingface_tasks/computer_vision/zero_shot_image_classification/node.py new file mode 100644 index 0000000..5bef890 --- /dev/null +++ b/huggingface_tasks/computer_vision/zero_shot_image_classification/node.py @@ -0,0 +1,35 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np +import json + +zero_shot_image_classification_model_list = [ + "openai/clip-vit-base-patch32", + "openai/clip-vit-large-patch14", +] + + +class ZeroShotImageClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "candidate_labels": ("STRING", {"default": "cat, dog, bird"}), + "model_name": (zero_shot_image_classification_model_list, {"default": zero_shot_image_classification_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("results_json",) + FUNCTION = "run_zero_shot_classification" + CATEGORY = "Transformers/ComputerVision/ZeroShotImageClassification" + + def run_zero_shot_classification(self, image, candidate_labels, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + labels = [l.strip() for l in candidate_labels.split(",")] + pipe = hf_pipeline("zero-shot-image-classification", model=model_name) + result = pipe(pil_image, candidate_labels=labels) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/computer_vision/zero_shot_object_detection/node.py b/huggingface_tasks/computer_vision/zero_shot_object_detection/node.py new file mode 100644 index 0000000..f049e47 --- /dev/null +++ b/huggingface_tasks/computer_vision/zero_shot_object_detection/node.py @@ -0,0 +1,48 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image, ImageDraw +import numpy as np +import torch +import json + +zero_shot_object_detection_model_list = [ + "google/owlvit-base-patch32", + "google/owlv2-base-patch16-ensemble", +] + + +class ZeroShotObjectDetectionPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "candidate_labels": ("STRING", {"default": "cat, dog, person"}), + "model_name": (zero_shot_object_detection_model_list, {"default": zero_shot_object_detection_model_list[0]}), + "threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("annotated_image", "detections_json",) + FUNCTION = "run_zero_shot_detection" + CATEGORY = "Transformers/ComputerVision/ZeroShotObjectDetection" + + def run_zero_shot_detection(self, image, candidate_labels, model_name, threshold): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + labels = [l.strip() for l in candidate_labels.split(",")] + pipe = hf_pipeline("zero-shot-object-detection", model=model_name) + results = pipe(pil_image, candidate_labels=labels, threshold=threshold) + + draw = ImageDraw.Draw(pil_image) + for det in results: + box = det["box"] + x, y, x2, y2 = box["xmin"], box["ymin"], box["xmax"], box["ymax"] + draw.rectangle((x, y, x2, y2), outline="red", width=2) + label_text = f"{det['label']} ({det['score']:.2f})" + draw.text((x, y), label_text, fill="white") + + arr = np.array(pil_image).astype(np.float32) / 255.0 + tensor_image = torch.from_numpy(arr)[None,] + return (tensor_image, json.dumps(results, indent=2, default=str),) diff --git a/huggingface_tasks/multimodal/document_question_answering/node.py b/huggingface_tasks/multimodal/document_question_answering/node.py index e69de29..5aa73c1 100644 --- a/huggingface_tasks/multimodal/document_question_answering/node.py +++ b/huggingface_tasks/multimodal/document_question_answering/node.py @@ -0,0 +1,34 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np + +doc_qa_model_list = [ + "impira/layoutlm-document-qa", + "naver-clova-ix/donut-base-finetuned-docvqa", +] + + +class DocumentQuestionAnsweringPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "question": ("STRING", {"default": "", "multiline": False}), + "model_name": (doc_qa_model_list, {"default": doc_qa_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING", "FLOAT",) + RETURN_NAMES = ("answer", "score",) + FUNCTION = "run_doc_qa" + CATEGORY = "Transformers/Multimodal/DocumentQA" + + def run_doc_qa(self, image, question, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("document-question-answering", model=model_name) + result = pipe(pil_image, question) + top = result[0] + return (top["answer"], top.get("score", 0.0),) diff --git a/huggingface_tasks/multimodal/feature_extraction/node.py b/huggingface_tasks/multimodal/feature_extraction/node.py index e69de29..c1c96ca 100644 --- a/huggingface_tasks/multimodal/feature_extraction/node.py +++ b/huggingface_tasks/multimodal/feature_extraction/node.py @@ -0,0 +1,31 @@ +from transformers import pipeline as hf_pipeline +import numpy as np +import json + +feature_extraction_model_list = [ + "sentence-transformers/all-MiniLM-L6-v2", + "sentence-transformers/all-mpnet-base-v2", +] + + +class FeatureExtractionPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (feature_extraction_model_list, {"default": feature_extraction_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("features_json",) + FUNCTION = "run_feature_extraction" + CATEGORY = "Transformers/Multimodal/FeatureExtraction" + + def run_feature_extraction(self, text, model_name): + pipe = hf_pipeline("feature-extraction", model=model_name) + result = pipe(text) + features = np.array(result[0]) + output = {"shape": list(features.shape), "features_preview": features[0][:10].tolist()} + return (json.dumps(output, indent=2),) diff --git a/huggingface_tasks/multimodal/image2text/node.py b/huggingface_tasks/multimodal/image2text/node.py index e69de29..14a5b13 100644 --- a/huggingface_tasks/multimodal/image2text/node.py +++ b/huggingface_tasks/multimodal/image2text/node.py @@ -0,0 +1,33 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np + +image2text_model_list = [ + "Salesforce/blip-image-captioning-base", + "Salesforce/blip-image-captioning-large", +] + + +class ImageToTextPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": (image2text_model_list, {"default": image2text_model_list[0]}), + "max_new_tokens": ("INT", {"default": 50, "min": 1, "max": 512}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("generated_text",) + FUNCTION = "run_image2text" + CATEGORY = "Transformers/Multimodal/ImageToText" + + def run_image2text(self, image, model_name, max_new_tokens): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("image-to-text", model=model_name) + result = pipe(pil_image, max_new_tokens=max_new_tokens) + return (result[0]["generated_text"],) diff --git a/huggingface_tasks/multimodal/image_text_to_text/node.py b/huggingface_tasks/multimodal/image_text_to_text/node.py new file mode 100644 index 0000000..865eb4a --- /dev/null +++ b/huggingface_tasks/multimodal/image_text_to_text/node.py @@ -0,0 +1,38 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np + +image_text_to_text_model_list = [ + "Salesforce/blip2-opt-2.7b", + "llava-hf/llava-1.5-7b-hf", +] + + +class ImageTextToTextPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"default": "", "multiline": True}), + "model_name": (image_text_to_text_model_list, {"default": image_text_to_text_model_list[0]}), + "max_new_tokens": ("INT", {"default": 50, "min": 1, "max": 512}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("generated_text",) + FUNCTION = "run_image_text_to_text" + CATEGORY = "Transformers/Multimodal/ImageTextToText" + + def run_image_text_to_text(self, image, prompt, model_name, max_new_tokens): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("image-text-to-text", model=model_name) + messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}] + result = pipe(pil_image, text=messages, max_new_tokens=max_new_tokens) + if isinstance(result, list) and len(result) > 0: + if isinstance(result[0], dict): + return (result[0].get("generated_text", str(result[0])),) + return (str(result),) diff --git a/huggingface_tasks/multimodal/visual_question_answering/node.py b/huggingface_tasks/multimodal/visual_question_answering/node.py index e69de29..c89c222 100644 --- a/huggingface_tasks/multimodal/visual_question_answering/node.py +++ b/huggingface_tasks/multimodal/visual_question_answering/node.py @@ -0,0 +1,34 @@ +from transformers import pipeline as hf_pipeline +from PIL import Image +import numpy as np +import json + +vqa_model_list = [ + "dandelin/vilt-b32-finetuned-vqa", + "Salesforce/blip-vqa-base", +] + + +class VisualQuestionAnsweringPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "question": ("STRING", {"default": "", "multiline": False}), + "model_name": (vqa_model_list, {"default": vqa_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("answers_json",) + FUNCTION = "run_vqa" + CATEGORY = "Transformers/Multimodal/VisualQA" + + def run_vqa(self, image, question, model_name): + img = 255.0 * image[0].cpu().numpy() + pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + pipe = hf_pipeline("visual-question-answering", model=model_name) + result = pipe(pil_image, question) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/natural_language_processing/conversational/node.py b/huggingface_tasks/natural_language_processing/conversational/node.py index e69de29..18c1954 100644 --- a/huggingface_tasks/natural_language_processing/conversational/node.py +++ b/huggingface_tasks/natural_language_processing/conversational/node.py @@ -0,0 +1,40 @@ +from transformers import pipeline as hf_pipeline +import json + +conversational_model_list = [ + "microsoft/DialoGPT-medium", + "microsoft/DialoGPT-small", +] + + +class ConversationalPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "user_message": ("STRING", {"default": "", "multiline": True}), + "history": ("STRING", {"default": "[]", "multiline": True}), + "model_name": (conversational_model_list, {"default": conversational_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING", "STRING",) + RETURN_NAMES = ("response", "updated_history",) + FUNCTION = "run_conversation" + CATEGORY = "Transformers/NLP/Conversational" + + def run_conversation(self, user_message, history, model_name): + history_list = json.loads(history) if history.strip() else [] + + messages = [] + for turn in history_list: + messages.append({"role": "user", "content": turn["user"]}) + messages.append({"role": "assistant", "content": turn["assistant"]}) + messages.append({"role": "user", "content": user_message}) + + pipe = hf_pipeline("text-generation", model=model_name) + result = pipe(messages, max_new_tokens=128, pad_token_id=pipe.tokenizer.eos_token_id) + response = result[0]["generated_text"][-1]["content"] + + history_list.append({"user": user_message, "assistant": response}) + return (response, json.dumps(history_list, indent=2),) diff --git a/huggingface_tasks/natural_language_processing/fill_mask/node.py b/huggingface_tasks/natural_language_processing/fill_mask/node.py index e69de29..299c015 100644 --- a/huggingface_tasks/natural_language_processing/fill_mask/node.py +++ b/huggingface_tasks/natural_language_processing/fill_mask/node.py @@ -0,0 +1,29 @@ +from transformers import pipeline as hf_pipeline +import json + +fill_mask_model_list = [ + "bert-base-uncased", + "distilbert-base-uncased", +] + + +class FillMaskPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "The capital of France is [MASK].", "multiline": True}), + "model_name": (fill_mask_model_list, {"default": fill_mask_model_list[0]}), + "top_k": ("INT", {"default": 5, "min": 1, "max": 20}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("predictions_json",) + FUNCTION = "run_fill_mask" + CATEGORY = "Transformers/NLP/FillMask" + + def run_fill_mask(self, text, model_name, top_k): + pipe = hf_pipeline("fill-mask", model=model_name) + result = pipe(text, top_k=top_k) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/natural_language_processing/question_answering/node.py b/huggingface_tasks/natural_language_processing/question_answering/node.py index e69de29..25faa59 100644 --- a/huggingface_tasks/natural_language_processing/question_answering/node.py +++ b/huggingface_tasks/natural_language_processing/question_answering/node.py @@ -0,0 +1,28 @@ +from transformers import pipeline as hf_pipeline + +question_answering_model_list = [ + "distilbert-base-cased-distilled-squad", + "deepset/roberta-base-squad2", +] + + +class QuestionAnsweringPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "question": ("STRING", {"default": "", "multiline": False}), + "context": ("STRING", {"default": "", "multiline": True}), + "model_name": (question_answering_model_list, {"default": question_answering_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING", "FLOAT",) + RETURN_NAMES = ("answer", "score",) + FUNCTION = "run_qa" + CATEGORY = "Transformers/NLP/QuestionAnswering" + + def run_qa(self, question, context, model_name): + pipe = hf_pipeline("question-answering", model=model_name) + result = pipe(question=question, context=context) + return (result["answer"], result["score"],) diff --git a/huggingface_tasks/natural_language_processing/sentence_similarity/node.py b/huggingface_tasks/natural_language_processing/sentence_similarity/node.py index e69de29..bed5751 100644 --- a/huggingface_tasks/natural_language_processing/sentence_similarity/node.py +++ b/huggingface_tasks/natural_language_processing/sentence_similarity/node.py @@ -0,0 +1,32 @@ +from transformers import pipeline as hf_pipeline +import numpy as np + +sentence_similarity_model_list = [ + "sentence-transformers/all-MiniLM-L6-v2", + "sentence-transformers/all-mpnet-base-v2", +] + + +class SentenceSimilarityPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text1": ("STRING", {"default": "", "multiline": True}), + "text2": ("STRING", {"default": "", "multiline": True}), + "model_name": (sentence_similarity_model_list, {"default": sentence_similarity_model_list[0]}), + }, + } + + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("similarity_score",) + FUNCTION = "run_similarity" + CATEGORY = "Transformers/NLP/SentenceSimilarity" + + def run_similarity(self, text1, text2, model_name): + pipe = hf_pipeline("feature-extraction", model=model_name) + emb1 = np.array(pipe(text1)[0]).mean(axis=0) + emb2 = np.array(pipe(text2)[0]).mean(axis=0) + + cosine_sim = float(np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))) + return (cosine_sim,) diff --git a/huggingface_tasks/natural_language_processing/summarization/node.py b/huggingface_tasks/natural_language_processing/summarization/node.py index e69de29..551c551 100644 --- a/huggingface_tasks/natural_language_processing/summarization/node.py +++ b/huggingface_tasks/natural_language_processing/summarization/node.py @@ -0,0 +1,29 @@ +from transformers import pipeline as hf_pipeline + +summarization_model_list = [ + "facebook/bart-large-cnn", + "sshleifer/distilbart-cnn-12-6", +] + + +class SummarizationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (summarization_model_list, {"default": summarization_model_list[0]}), + "max_length": ("INT", {"default": 130, "min": 10, "max": 1024}), + "min_length": ("INT", {"default": 30, "min": 1, "max": 512}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("summary_text",) + FUNCTION = "run_summarization" + CATEGORY = "Transformers/NLP/Summarization" + + def run_summarization(self, text, model_name, max_length, min_length): + pipe = hf_pipeline("summarization", model=model_name) + result = pipe(text, max_length=max_length, min_length=min_length) + return (result[0]["summary_text"],) diff --git a/huggingface_tasks/natural_language_processing/table_question_answering/node.py b/huggingface_tasks/natural_language_processing/table_question_answering/node.py index e69de29..cee0f98 100644 --- a/huggingface_tasks/natural_language_processing/table_question_answering/node.py +++ b/huggingface_tasks/natural_language_processing/table_question_answering/node.py @@ -0,0 +1,30 @@ +from transformers import pipeline as hf_pipeline +import json + +table_qa_model_list = [ + "google/tapas-base-finetuned-wtq", + "google/tapas-large-finetuned-wtq", +] + + +class TableQuestionAnsweringPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "question": ("STRING", {"default": "", "multiline": False}), + "table_json": ("STRING", {"default": '{"Name": ["Alice", "Bob"], "Age": ["25", "30"]}', "multiline": True}), + "model_name": (table_qa_model_list, {"default": table_qa_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("answer",) + FUNCTION = "run_table_qa" + CATEGORY = "Transformers/NLP/TableQuestionAnswering" + + def run_table_qa(self, question, table_json, model_name): + table = json.loads(table_json) + pipe = hf_pipeline("table-question-answering", model=model_name) + result = pipe(table=table, query=question) + return (result["answer"],) diff --git a/huggingface_tasks/natural_language_processing/text2text_generation/node.py b/huggingface_tasks/natural_language_processing/text2text_generation/node.py new file mode 100644 index 0000000..49246ea --- /dev/null +++ b/huggingface_tasks/natural_language_processing/text2text_generation/node.py @@ -0,0 +1,28 @@ +from transformers import pipeline as hf_pipeline + +text2text_generation_model_list = [ + "google/flan-t5-small", + "google/flan-t5-base", +] + + +class Text2TextGenerationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (text2text_generation_model_list, {"default": text2text_generation_model_list[0]}), + "max_new_tokens": ("INT", {"default": 50, "min": 1, "max": 2048}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("generated_text",) + FUNCTION = "run_text2text" + CATEGORY = "Transformers/NLP/Text2TextGeneration" + + def run_text2text(self, text, model_name, max_new_tokens): + pipe = hf_pipeline("text2text-generation", model=model_name) + result = pipe(text, max_new_tokens=max_new_tokens) + return (result[0]["generated_text"],) diff --git a/huggingface_tasks/natural_language_processing/text_classification/node.py b/huggingface_tasks/natural_language_processing/text_classification/node.py index e69de29..a4ab960 100644 --- a/huggingface_tasks/natural_language_processing/text_classification/node.py +++ b/huggingface_tasks/natural_language_processing/text_classification/node.py @@ -0,0 +1,28 @@ +from transformers import pipeline as hf_pipeline + +text_classification_model_list = [ + "distilbert-base-uncased-finetuned-sst-2-english", + "cardiffnlp/twitter-roberta-base-sentiment-latest", +] + + +class TextClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (text_classification_model_list, {"default": text_classification_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING", "FLOAT",) + RETURN_NAMES = ("label", "score",) + FUNCTION = "run_classification" + CATEGORY = "Transformers/NLP/TextClassification" + + def run_classification(self, text, model_name): + pipe = hf_pipeline("text-classification", model=model_name) + result = pipe(text) + top = result[0] + return (top["label"], top["score"],) diff --git a/huggingface_tasks/natural_language_processing/text_generation/node.py b/huggingface_tasks/natural_language_processing/text_generation/node.py index e69de29..0befee5 100644 --- a/huggingface_tasks/natural_language_processing/text_generation/node.py +++ b/huggingface_tasks/natural_language_processing/text_generation/node.py @@ -0,0 +1,29 @@ +from transformers import pipeline as hf_pipeline + +text_generation_model_list = [ + "gpt2", + "distilgpt2", +] + + +class TextGenerationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"default": "", "multiline": True}), + "model_name": (text_generation_model_list, {"default": text_generation_model_list[0]}), + "max_new_tokens": ("INT", {"default": 50, "min": 1, "max": 2048}), + "temperature": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 2.0, "step": 0.05}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("generated_text",) + FUNCTION = "run_generation" + CATEGORY = "Transformers/NLP/TextGeneration" + + def run_generation(self, prompt, model_name, max_new_tokens, temperature): + pipe = hf_pipeline("text-generation", model=model_name) + result = pipe(prompt, max_new_tokens=max_new_tokens, temperature=temperature) + return (result[0]["generated_text"],) diff --git a/huggingface_tasks/natural_language_processing/token_classification/node.py b/huggingface_tasks/natural_language_processing/token_classification/node.py index e69de29..815a0bb 100644 --- a/huggingface_tasks/natural_language_processing/token_classification/node.py +++ b/huggingface_tasks/natural_language_processing/token_classification/node.py @@ -0,0 +1,37 @@ +from transformers import pipeline as hf_pipeline +import json + +token_classification_model_list = [ + "dslim/bert-base-NER", + "Jean-Baptiste/camembert-ner", +] + +aggregation_strategy_list = [ + "simple", + "first", + "average", + "max", + "none", +] + + +class TokenClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (token_classification_model_list, {"default": token_classification_model_list[0]}), + "aggregation_strategy": (aggregation_strategy_list, {"default": "simple"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("entities_json",) + FUNCTION = "run_token_classification" + CATEGORY = "Transformers/NLP/TokenClassification" + + def run_token_classification(self, text, model_name, aggregation_strategy): + pipe = hf_pipeline("token-classification", model=model_name, aggregation_strategy=aggregation_strategy) + result = pipe(text) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/natural_language_processing/translation/node.py b/huggingface_tasks/natural_language_processing/translation/node.py index e69de29..ec5326b 100644 --- a/huggingface_tasks/natural_language_processing/translation/node.py +++ b/huggingface_tasks/natural_language_processing/translation/node.py @@ -0,0 +1,28 @@ +from transformers import pipeline as hf_pipeline + +translation_model_list = [ + "Helsinki-NLP/opus-mt-en-de", + "Helsinki-NLP/opus-mt-en-fr", +] + + +class TranslationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "model_name": (translation_model_list, {"default": translation_model_list[0]}), + "max_length": ("INT", {"default": 512, "min": 10, "max": 2048}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("translation_text",) + FUNCTION = "run_translation" + CATEGORY = "Transformers/NLP/Translation" + + def run_translation(self, text, model_name, max_length): + pipe = hf_pipeline("translation", model=model_name) + result = pipe(text, max_length=max_length) + return (result[0]["translation_text"],) diff --git a/huggingface_tasks/natural_language_processing/zero_shot_classification/node.py b/huggingface_tasks/natural_language_processing/zero_shot_classification/node.py index e69de29..a107a32 100644 --- a/huggingface_tasks/natural_language_processing/zero_shot_classification/node.py +++ b/huggingface_tasks/natural_language_processing/zero_shot_classification/node.py @@ -0,0 +1,30 @@ +from transformers import pipeline as hf_pipeline +import json + +zero_shot_classification_model_list = [ + "facebook/bart-large-mnli", + "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli", +] + + +class ZeroShotClassificationPipeline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"default": "", "multiline": True}), + "candidate_labels": ("STRING", {"default": "positive, negative, neutral"}), + "model_name": (zero_shot_classification_model_list, {"default": zero_shot_classification_model_list[0]}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("results_json",) + FUNCTION = "run_zero_shot" + CATEGORY = "Transformers/NLP/ZeroShotClassification" + + def run_zero_shot(self, text, candidate_labels, model_name): + labels = [l.strip() for l in candidate_labels.split(",")] + pipe = hf_pipeline("zero-shot-classification", model=model_name) + result = pipe(text, candidate_labels=labels) + return (json.dumps(result, indent=2, default=str),) diff --git a/huggingface_tasks/utils.py b/huggingface_tasks/utils.py new file mode 100644 index 0000000..139070f --- /dev/null +++ b/huggingface_tasks/utils.py @@ -0,0 +1,36 @@ +import numpy as np +import torch +from PIL import Image +from transformers import pipeline as hf_pipeline +import json + +_pipeline_cache = {} + + +def get_pipeline(task, model_name, **kwargs): + """Cached pipeline loader - avoids reloading models on every call.""" + key = (task, model_name, frozenset(kwargs.items())) + if key not in _pipeline_cache: + _pipeline_cache[key] = hf_pipeline(task, model=model_name, **kwargs) + return _pipeline_cache[key] + + +def comfyui_to_pil(image_tensor): + """ComfyUI IMAGE tensor (B,H,W,C) [0,1] -> PIL Image (first frame).""" + img = 255.0 * image_tensor[0].cpu().numpy() + return Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + + +def pil_to_comfyui(pil_image): + """PIL Image -> ComfyUI IMAGE tensor (1,H,W,C) [0,1].""" + arr = np.array(pil_image).astype(np.float32) / 255.0 + if len(arr.shape) == 2: + arr = np.stack([arr, arr, arr], axis=-1) + elif arr.shape[-1] == 4: + arr = arr[:, :, :3] + return torch.from_numpy(arr)[None,] + + +def format_output(data): + """Convert pipeline output to formatted JSON string.""" + return json.dumps(data, indent=2, default=str) diff --git a/nodes.py b/nodes.py index e378bf7..2c3b062 100644 --- a/nodes.py +++ b/nodes.py @@ -1,21 +1,120 @@ +# Computer Vision from .huggingface_tasks.computer_vision.depth_estimation.node import LoadDepthModel, DepthEstimationInference from .huggingface_tasks.computer_vision.image_classification.node import ImageClassificationPipeline from .huggingface_tasks.computer_vision.image_segmentation.node import ImageSegmentationPipeline from .huggingface_tasks.computer_vision.object_detection.node import ObjectDetectionPipeline +from .huggingface_tasks.computer_vision.image_to_image.node import ImageToImagePipeline +from .huggingface_tasks.computer_vision.image_feature_extraction.node import ImageFeatureExtractionPipeline +from .huggingface_tasks.computer_vision.zero_shot_image_classification.node import ZeroShotImageClassificationPipeline +from .huggingface_tasks.computer_vision.zero_shot_object_detection.node import ZeroShotObjectDetectionPipeline +from .huggingface_tasks.computer_vision.mask_generation.node import MaskGenerationPipeline +from .huggingface_tasks.computer_vision.video_classification.node import VideoClassificationPipeline +from .huggingface_tasks.computer_vision.keypoint_detection.node import KeypointDetectionPipeline + +# NLP +from .huggingface_tasks.natural_language_processing.text_classification.node import TextClassificationPipeline +from .huggingface_tasks.natural_language_processing.text_generation.node import TextGenerationPipeline +from .huggingface_tasks.natural_language_processing.text2text_generation.node import Text2TextGenerationPipeline +from .huggingface_tasks.natural_language_processing.summarization.node import SummarizationPipeline +from .huggingface_tasks.natural_language_processing.translation.node import TranslationPipeline +from .huggingface_tasks.natural_language_processing.token_classification.node import TokenClassificationPipeline +from .huggingface_tasks.natural_language_processing.fill_mask.node import FillMaskPipeline +from .huggingface_tasks.natural_language_processing.question_answering.node import QuestionAnsweringPipeline +from .huggingface_tasks.natural_language_processing.zero_shot_classification.node import ZeroShotClassificationPipeline +from .huggingface_tasks.natural_language_processing.table_question_answering.node import TableQuestionAnsweringPipeline +from .huggingface_tasks.natural_language_processing.conversational.node import ConversationalPipeline +from .huggingface_tasks.natural_language_processing.sentence_similarity.node import SentenceSimilarityPipeline + +# Audio +from .huggingface_tasks.audio.automatic_speech_recognition.node import AutomaticSpeechRecognitionPipeline +from .huggingface_tasks.audio.audio_classification.node import AudioClassificationPipeline +from .huggingface_tasks.audio.text2speech.node import TextToSpeechPipeline +from .huggingface_tasks.audio.zero_shot_audio_classification.node import ZeroShotAudioClassificationPipeline + +# Multimodal +from .huggingface_tasks.multimodal.feature_extraction.node import FeatureExtractionPipeline +from .huggingface_tasks.multimodal.image2text.node import ImageToTextPipeline +from .huggingface_tasks.multimodal.visual_question_answering.node import VisualQuestionAnsweringPipeline +from .huggingface_tasks.multimodal.document_question_answering.node import DocumentQuestionAnsweringPipeline +from .huggingface_tasks.multimodal.image_text_to_text.node import ImageTextToTextPipeline NODE_CLASS_MAPPINGS = { - "LoadDepthModel": LoadDepthModel, - "DepthEstimationInference": DepthEstimationInference, - "ImageClassificationPipeline": ImageClassificationPipeline, - "ImageSegmentationPipeline": ImageSegmentationPipeline, - "ObjectDetectionPipeline": ObjectDetectionPipeline, + # Computer Vision + "LoadDepthModel": LoadDepthModel, + "DepthEstimationInference": DepthEstimationInference, + "ImageClassificationPipeline": ImageClassificationPipeline, + "ImageSegmentationPipeline": ImageSegmentationPipeline, + "ObjectDetectionPipeline": ObjectDetectionPipeline, + "ImageToImagePipeline": ImageToImagePipeline, + "ImageFeatureExtractionPipeline": ImageFeatureExtractionPipeline, + "ZeroShotImageClassificationPipeline": ZeroShotImageClassificationPipeline, + "ZeroShotObjectDetectionPipeline": ZeroShotObjectDetectionPipeline, + "MaskGenerationPipeline": MaskGenerationPipeline, + "VideoClassificationPipeline": VideoClassificationPipeline, + "KeypointDetectionPipeline": KeypointDetectionPipeline, + # NLP + "TextClassificationPipeline": TextClassificationPipeline, + "TextGenerationPipeline": TextGenerationPipeline, + "Text2TextGenerationPipeline": Text2TextGenerationPipeline, + "SummarizationPipeline": SummarizationPipeline, + "TranslationPipeline": TranslationPipeline, + "TokenClassificationPipeline": TokenClassificationPipeline, + "FillMaskPipeline": FillMaskPipeline, + "QuestionAnsweringPipeline": QuestionAnsweringPipeline, + "ZeroShotClassificationPipeline": ZeroShotClassificationPipeline, + "TableQuestionAnsweringPipeline": TableQuestionAnsweringPipeline, + "ConversationalPipeline": ConversationalPipeline, + "SentenceSimilarityPipeline": SentenceSimilarityPipeline, + # Audio + "AutomaticSpeechRecognitionPipeline": AutomaticSpeechRecognitionPipeline, + "AudioClassificationPipeline": AudioClassificationPipeline, + "TextToSpeechPipeline": TextToSpeechPipeline, + "ZeroShotAudioClassificationPipeline": ZeroShotAudioClassificationPipeline, + # Multimodal + "FeatureExtractionPipeline": FeatureExtractionPipeline, + "ImageToTextPipeline": ImageToTextPipeline, + "VisualQuestionAnsweringPipeline": VisualQuestionAnsweringPipeline, + "DocumentQuestionAnsweringPipeline": DocumentQuestionAnsweringPipeline, + "ImageTextToTextPipeline": ImageTextToTextPipeline, } NODE_DISPLAY_NAME_MAPPINGS = { - "LoadDepthModel": "LoadDepthModel", - "DepthEstimationInference": "DepthEstimation", - "ImageClassificationPipeline": "ImageClassification", - "ImageSegmentationPipeline": "ImageSegmentation", - "ObjectDetectionPipeline": "ObjectDetection", + # Computer Vision + "LoadDepthModel": "Load Depth Model", + "DepthEstimationInference": "Depth Estimation", + "ImageClassificationPipeline": "Image Classification", + "ImageSegmentationPipeline": "Image Segmentation", + "ObjectDetectionPipeline": "Object Detection", + "ImageToImagePipeline": "Image to Image", + "ImageFeatureExtractionPipeline": "Image Feature Extraction", + "ZeroShotImageClassificationPipeline": "Zero-Shot Image Classification", + "ZeroShotObjectDetectionPipeline": "Zero-Shot Object Detection", + "MaskGenerationPipeline": "Mask Generation", + "VideoClassificationPipeline": "Video Classification", + "KeypointDetectionPipeline": "Keypoint Detection", + # NLP + "TextClassificationPipeline": "Text Classification", + "TextGenerationPipeline": "Text Generation", + "Text2TextGenerationPipeline": "Text2Text Generation", + "SummarizationPipeline": "Summarization", + "TranslationPipeline": "Translation", + "TokenClassificationPipeline": "Token Classification (NER)", + "FillMaskPipeline": "Fill Mask", + "QuestionAnsweringPipeline": "Question Answering", + "ZeroShotClassificationPipeline": "Zero-Shot Classification", + "TableQuestionAnsweringPipeline": "Table Question Answering", + "ConversationalPipeline": "Conversational", + "SentenceSimilarityPipeline": "Sentence Similarity", + # Audio + "AutomaticSpeechRecognitionPipeline": "Speech Recognition (ASR)", + "AudioClassificationPipeline": "Audio Classification", + "TextToSpeechPipeline": "Text to Speech", + "ZeroShotAudioClassificationPipeline": "Zero-Shot Audio Classification", + # Multimodal + "FeatureExtractionPipeline": "Feature Extraction", + "ImageToTextPipeline": "Image to Text", + "VisualQuestionAnsweringPipeline": "Visual Question Answering", + "DocumentQuestionAnsweringPipeline": "Document Question Answering", + "ImageTextToTextPipeline": "Image-Text to Text", } diff --git a/requirements.txt b/requirements.txt index 7e85007..99d19d2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,9 @@ -numpy=>1.26.3 -Pillow=>10.2.0 -torch=>2.1.2 -transformers=>4.37.1 +numpy>=1.26.3 +Pillow>=10.2.0 +torch>=2.1.2 +transformers>=4.42.0 +scipy>=1.12.0 +soundfile>=0.12.1 +pandas>=2.1.0 +sentencepiece>=0.1.99 +accelerate>=0.25.0