From 355a80b4c4f9a57b5c7c76e9af0d7caab24eb7f7 Mon Sep 17 00:00:00 2001
From: lining3 <lining3@onething.net>
Date: Tue, 15 Jul 2025 14:42:06 +0800
Subject: [PATCH] add reference image support

---
 README.md     | 91 ++++++++++++++++++++++++++-------------------------
 onethingai.py | 28 +++++++++++++++-
 2 files changed, 74 insertions(+), 45 deletions(-)

diff --git a/README.md b/README.md
index 7ff8c6e..94622d4 100644
--- a/README.md
+++ b/README.md
@@ -1,44 +1,47 @@
-# ComfyUI OneThing AI Node
-
-A custom node for ComfyUI that integrates with OneThing AI's image generation API.
-
-## Features
-- Generate images using OneThing AI's models
-- Control image dimensions (512x512 to 2048x2048)
-- Multiple images per request (1-10)
-- Built-in retry mechanism for API stability
-- Configurable timeout settings
-
-## Installation
-
-1. Clone this repository into your ComfyUI custom_nodes folder:
-```bash
-cd ComfyUI/custom_nodes
-git clone https://github.com/OneThingAI/ComfyUI_Onething_Image.git
-```
-
-2. Install required dependencies:
-```bash
-pip install Pillow requests
-```
-
-## Usage
-
-1. Get your API key from OneThing AI
-2. In ComfyUI, find the "OneThingAI Image Generator" node
-3. Configure the parameters:
-   - API Key: Your OneThing AI API key
-   - Prompt: Text description of the image you want to generate
-   - Model: Select the model (default: gpt4o)
-   - Number of Images: How many images to generate (1-10)
-   - Width/Height: Image dimensions (512-2048, step 64)
-   - Retries: Number of retry attempts (0-5)
-   - Timeout: Request timeout in seconds (5-100)
-
-## Requirements
-- Python 3.x
-- ComfyUI
-- Pillow (PIL)
-- requests
-
-
+# ComfyUI OneThing AI Node
+
+A custom node for ComfyUI that integrates with OneThing AI's image generation API.
+
+## Features
+- Generate images using OneThing AI's models
+- Control image dimensions (512x512 to 2048x2048)
+- Multiple images per request (1-10)
+- Built-in retry mechanism for API stability
+- Configurable timeout settings
+- Reference image support for guided image generation
+
+## Installation
+
+1. Clone this repository into your ComfyUI custom_nodes folder:
+```bash
+cd ComfyUI/custom_nodes
+git clone https://github.com/OneThingAI/ComfyUI_Onething_Image.git
+```
+
+2. Install required dependencies:
+```bash
+pip install Pillow requests
+```
+
+## Usage
+
+1. Get your API key from OneThing AI
+2. In ComfyUI, find the "OneThingAI Image Generator" node
+3. Configure the parameters:
+   - API Key: Your OneThing AI API key
+   - Prompt: Text description of the image you want to generate
+   - Model: Select the model (default: gpt4o)
+   - Number of Images: How many images to generate (1-10)
+   - Width/Height: Image dimensions (512-2048, step 64)
+   - Retries: Number of retry attempts (0-5)
+   - Timeout: Request timeout in seconds (5-100)
+   - Reference Image (optional): Input image to guide the generation
+   - Reference Image Weight (optional): Control how much the reference image influences the result (0.0-1.0)
+
+## Requirements
+- Python 3.x
+- ComfyUI
+- Pillow (PIL)
+- requests
+
+
diff --git a/onethingai.py b/onethingai.py
index f3333ef..44ce96a 100644
--- a/onethingai.py
+++ b/onethingai.py
@@ -22,13 +22,17 @@ class OneThingAILoader:
                 "retries": ("INT", {"default": 3, "min": 0, "max": 5}),
                 "timeout": ("INT", {"default": 20, "min": 5, "max": 100}),
             },
+            "optional": {
+                "reference_image": ("IMAGE",),
+                "reference_image_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
+            }
         }
     
     RETURN_TYPES = ("IMAGE", )
     FUNCTION = "generate"
     CATEGORY = "OneThingAI/image generation"
 
-    def generate(self, api_key, prompt, model="gpt4o", num_images=1, width=1024, height=1024, retries=3, timeout=8):
+    def generate(self, api_key, prompt, model="gpt4o", num_images=1, width=1024, height=1024, retries=3, timeout=8, reference_image=None, reference_image_weight=0.5):
         # API endpoint
         url = "https://api-model.onethingai.com/v1/images/generations"
         
@@ -57,6 +61,28 @@ class OneThingAILoader:
             "n": num_images,
             "size": f"{width}x{height}",
         }
+
+        # Add reference image if provided
+        if reference_image is not None:
+            # Convert the first image if it's a batch
+            if len(reference_image.shape) == 4:
+                image_to_encode = reference_image[0]
+            else:
+                image_to_encode = reference_image
+
+            # Convert from torch tensor to PIL Image
+            image_to_encode = (image_to_encode.cpu().numpy() * 255).astype(np.uint8)
+            if len(image_to_encode.shape) == 3:
+                image_to_encode = Image.fromarray(image_to_encode)
+            
+            # Convert to base64
+            buffered = io.BytesIO()
+            image_to_encode.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            
+            # Add to payload
+            payload["reference_image"] = img_str
+            payload["reference_image_weight"] = reference_image_weight
             
         try:
             # Make API request with timeout