From fbf6e51eac75348e54fee6dea4a36ab2d17ec5f4 Mon Sep 17 00:00:00 2001
From: bitaffinity <168344121+bitaffinity@users.noreply.github.com>
Date: Fri, 31 May 2024 20:39:47 -0400
Subject: [PATCH] Add more to readme, Image nodes, and requirements.txt

---
 Image/nodes.py   | 73 ++++++++++++++++++++++++++++++++++++++++++++++++
 README.md        | 31 ++++++++++++++++++--
 Text/nodes.py    | 29 ++++++++++++++++++-
 __init__.py      |  3 ++
 requirements.txt |  1 +
 5 files changed, 133 insertions(+), 4 deletions(-)
 create mode 100644 Image/nodes.py
 create mode 100644 requirements.txt

diff --git a/Image/nodes.py b/Image/nodes.py
new file mode 100644
index 0000000..6b5c000
--- /dev/null
+++ b/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,
+}
\ No newline at end of file
diff --git a/README.md b/README.md
index 50f4cb4..36c02b6 100644
--- a/README.md
+++ b/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)
\ No newline at end of file
diff --git a/Text/nodes.py b/Text/nodes.py
index 5c1f2f4..0f7938c 100644
--- a/Text/nodes.py
+++ b/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,
 }
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
index 422976b..f9341e4 100644
--- a/__init__.py
+++ b/__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)
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..663bd1f
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+requests
\ No newline at end of file