init

__init__.py

@@ -0,0 +1,153 @@
+import base64
+import io
+from PIL import Image
+import numpy as np
+from openai import OpenAI
+
+
+class ImageToLLMPrompt:
+    """
+    A custom node that uploads an image and sends it to an OpenAI-compatible API
+    with optional system and user prompts, returning the LLM response as text.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "base_url": ("STRING", {
+                    "default": "http://localhost:8080/v1",
+                    "multiline": False
+                }),
+                "api_key": ("STRING", {
+                    "default": "sk-test",
+                    "multiline": False
+                }),
+                "model": ("STRING", {
+                    "default": "gpt-4o",
+                    "multiline": False
+                }),
+            },
+            "optional": {
+                "system_prompt": ("STRING", {
+                    "default": "",
+                    "multiline": True
+                }),
+                "user_prompt": ("STRING", {
+                    "default": "",
+                    "multiline": True
+                }),
+                "temperature": ("FLOAT", {
+                    "default": 0.7,
+                    "min": 0.0,
+                    "max": 2.0,
+                    "step": 0.1
+                }),
+                "max_tokens": ("INT", {
+                    "default": 2048,
+                    "min": 1,
+                    "max": 32768,
+                    "step": 1
+                }),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text",)
+    FUNCTION = "process_image"
+    CATEGORY = "LLM"
+    OUTPUT_NODE = True
+    
+    def tensor_to_base64(self, image_tensor):
+        """Convert ComfyUI image tensor to base64 string"""
+        # ComfyUI images are in format [B, H, W, C] with values 0-1
+        image_np = (image_tensor[0].cpu().numpy() * 255).astype(np.uint8)
+        
+        # Convert to PIL Image
+        pil_image = Image.fromarray(image_np)
+        
+        # Convert to base64
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        base64_image = base64.b64encode(buffer.read()).decode('utf-8')
+        
+        return base64_image
+    
+    def process_image(self, image, base_url, api_key, model, 
+                     system_prompt="", user_prompt="", temperature=0.7, max_tokens=2048):
+        """Send image to LLM and return response"""
+        
+        # Convert image to base64
+        base64_image = self.tensor_to_base64(image)
+        
+        # Create OpenAI client
+        client = OpenAI(
+            base_url=base_url,
+            api_key=api_key,
+            timeout=180
+        )
+        
+        # Build messages array
+        messages = []
+        
+        # Add optional system prompt
+        if system_prompt and system_prompt.strip():
+            messages.append({
+                "role": "system",
+                "content": system_prompt
+            })
+        
+        # Build user message content
+        user_content = []
+        
+        # Add optional user prompt text
+        if user_prompt and user_prompt.strip():
+            user_content.append({
+                "type": "text",
+                "text": user_prompt
+            })
+        
+        # Add the image
+        user_content.append({
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/png;base64,{base64_image}"
+            }
+        })
+        
+        messages.append({
+            "role": "user",
+            "content": user_content
+        })
+        
+        # Make API call
+        try:
+            print(f"[ImageToLLMPrompt] Sending request to {base_url} with model {model}")
+            response = client.chat.completions.create(
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                messages=messages
+            )
+            
+            result = response.choices[0].message.content
+            print(f"[ImageToLLMPrompt] LLM Response received (length: {len(result)} chars)")
+            print(f"[ImageToLLMPrompt] Response content:\n{result}")
+            return {"ui": {"text": (result,)}, "result": (result,)}
+            
+        except Exception as e:
+            error_msg = f"LLM API call failed: {str(e)}"
+            print(f"[ImageToLLMPrompt] ERROR: {error_msg}")
+            return {"ui": {"text": (error_msg,)}, "result": (error_msg,)}
+
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "ImageToLLMPrompt": ImageToLLMPrompt
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageToLLMPrompt": "Image to LLM Prompt"
+}