Add new pipelines for multimodal and audio tasks, enhance NLP capabilities
This commit is contained in:
kadirnar
@@ -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),)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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),)
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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),)
|
||||
@@ -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,)
|
||||
|
||||
@@ -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,],)
|
||||
@@ -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),)
|
||||
@@ -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),)
|
||||
@@ -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,)
|
||||
|
||||
@@ -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),)
|
||||
@@ -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),)
|
||||
@@ -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),)
|
||||
@@ -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),)
|
||||
|
||||
@@ -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),)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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),)
|
||||
@@ -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),)
|
||||
|
||||
@@ -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),)
|
||||
|
||||
@@ -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),)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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"],)
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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),)
|
||||
|
||||
@@ -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"],)
|
||||
|
||||
@@ -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),)
|
||||
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
+9
-4
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user