Add more to readme, Image nodes, and requirements.txt

Image/nodes.py

@@ -0,0 +1,73 @@
+from ..session import session
+
+class Classification:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "endpoint": ("STRING", {}),
+                "image": ("IMAGE",),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "inference"
+    CATEGORY = "HF_Inference/Image/Classification"
+    TITLE = "HF Image Classification"
+
+    def inference(self, endpoint, image):
+        response = session.post(endpoint, data=image)
+        if response.status_code != 200:
+            raise Exception(response.text)
+        result = response.json()
+        return {"ui": {"text": result}}
+    
+class ObjectDetection:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "endpoint": ("STRING", {}),
+                "image": ("IMAGE",),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "inference"
+    CATEGORY = "HF_Inference/Image/ObjectDetection"
+    TITLE = "HF Image Object Detection"
+
+    def inference(self, endpoint, image):
+        response = session.post(endpoint, data=image)
+        if response.status_code != 200:
+            raise Exception(response.text)
+        result = response.json()
+        return {"ui": {"text": result}}
+    
+class Segmentation:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "endpoint": ("STRING", {}),
+                "image": ("IMAGE",),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "inference"
+    CATEGORY = "HF_Inference/Image/Segmentation"
+    TITLE = "HF Image Segmentation"
+
+    def inference(self, endpoint, image):
+        response = session.post(endpoint, data=image)
+        if response.status_code != 200:
+            raise Exception(response.text)
+        result = response.json()
+        return {"ui": {"text": result}}
+
+NODE_CLASS_MAPPINGS = {
+    "Classification": Classification,
+    "ObjectDetection": ObjectDetection,
+    "Segmentation": Segmentation,
+}

README.md

@@ -1,11 +1,36 @@
 # Hugging Face hosted inference nodes for ComfyUI
 
-Unofficial support for Hugging Face's hosted inference
+Unofficial ComfyUI nodes for Hugging Face's inference API
+
+Visit [the official docs](https://huggingface.co/docs/api-inference/detailed_parameters) for an overview of how the HF inference endpoints work
+
+Find models by task on the [official website](https://huggingface.co/tasks)
 
 ## Installation
 
+### Clone and install dependencies
 ```
-git clone https://github.com/bitaffinity/ComfyUI_HF_Inference
-pip install requests
+git clone https://github.com/bitaffinity/ComfyUI_HF_Inference custom_nodes/ComfyUI_HF_Inference
+cd custom_nodes/ComfyUI_HF_Inference
+pip install -r requirements.txt
 ```
 
+Export HF_AUTH_TOKEN with one of your [Hugging Face tokens](https://huggingface.co/settings/tokens)
+
+### Run ComfyUI
+`HF_AUTH_TOKEN=hf_1111111111111111111111111111111111 python main.py`
+
+## Usage
+
+### Text
+
+* Feature Extraction (ie: T5 encoder embeddings, BERT, etc.)
+* Question Answering (ie: roberta-base for QA)
+* Translation (ie: T5 Small)
+* Generation (ie: zephyr-7b-beta)
+
+### Image
+
+* Classification (ie: vit-base-patch16-224)
+* Object Detection (ie: detr-resnet-50)
+* Segmentation (ie: detr-resnet-50)

Text/nodes.py

@@ -1,6 +1,32 @@
 import torch
 from ..session import session
 
+class Generation:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "endpoint": ("STRING", {}),
+                "text": ("STRING", {"multiline": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "inference"
+    CATEGORY = "HF_Inference/Text/Generation"
+    TITLE = "HF Text Generation"
+
+    def inference(self, endpoint, text):
+        json = {
+            'inputs': text,
+        }
+        response = session.post(endpoint, json=json)
+        if response.status_code != 200:
+            raise Exception(response.text)
+        result = response.json()
+        generated = ''.join(x['generated_text'] for x in result)
+        return generated 
+
 class Translation:
     @classmethod
     def INPUT_TYPES(s):
@@ -84,7 +110,8 @@ class FeatureExtraction:
         return ([[cond, {}]],)
 
 NODE_CLASS_MAPPINGS = {
-    "TextFeatureExtraction": FeatureExtraction,
+    "FeatureExtraction": FeatureExtraction,
     "QuestionAnswering": QuestionAnswering,
     "Translation": Translation,
+    "Generation": Generation,
 }

__init__.py

@@ -8,6 +8,9 @@ else:
     from .Text.nodes import NODE_CLASS_MAPPINGS as Text_Nodes
     NODE_CLASS_MAPPINGS.update(Text_Nodes)
 
+    from .Image.nodes import NODE_CLASS_MAPPINGS as Image_Nodes
+    NODE_CLASS_MAPPINGS.update(Image_Nodes)
+
     NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}
     print(NODE_CLASS_MAPPINGS)
     print(NODE_DISPLAY_NAME_MAPPINGS)

requirements.txt

@@ -0,0 +1 @@
+requests