Upload comfyui-claude-node-final.py

comfyui-claude-node-final.py

@@ -0,0 +1,84 @@
+import os
+import json
+import anthropic
+from typing import Dict, Any
+
+class ClaudeCustomPrompt:
+    """
+    ComfyUI node that generates prompts using Claude API
+    """
+    
+    def __init__(self):
+        self.api_key = os.getenv("ANTHROPIC_API_KEY", "")
+        self.client = None
+    
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "system_prompt": ("STRING", {
+                    "default": "You are an AI art prompt generator. Reply only with a prompt for image generation, no explanations.", 
+                    "multiline": True
+                }),
+                "user_input": ("STRING", {
+                    "default": "Generate a prompt for: space cat", 
+                    "multiline": True
+                }),
+            },
+            "optional": {
+                "api_key": ("STRING", {
+                    "default": "", 
+                    "multiline": False
+                }),
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "generate_prompt"
+    CATEGORY = "prompt generation"
+    
+    def generate_prompt(self, system_prompt: str, user_input: str, api_key: str = "") -> tuple[str]:
+        # Use provided API key or fallback to environment variable
+        key_to_use = api_key if api_key else self.api_key
+        if not key_to_use:
+            raise ValueError("No API key provided. Please set ANTHROPIC_API_KEY environment variable or provide it as input.")
+            
+        # Initialize client if needed
+        if not self.client or api_key:
+            self.client = anthropic.Anthropic(api_key=key_to_use)
+        
+        try:
+            # Make API call to Claude
+            message = self.client.messages.create(
+                model="claude-3-opus-20240229",
+                max_tokens=1000,
+                temperature=0,
+                system=system_prompt,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": user_input
+                            }
+                        ]
+                    }
+                ]
+            )
+            
+            # Extract the generated prompt
+            generated_prompt = message.content[0].text.strip()
+            return (generated_prompt,)
+            
+        except Exception as e:
+            raise RuntimeError(f"Error generating prompt: {str(e)}")
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "ClaudeCustomPrompt": ClaudeCustomPrompt
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ClaudeCustomPrompt": "Claude Prompt Generator"
+}