Add files via upload

together_image_node.py

@@ -26,6 +26,25 @@ def __init__(self):
         self.last_request_time = 0
         self.min_request_interval = 1.0  # Minimum seconds between requests
 
+    def get_client(self, api_key: str) -> Optional[Together]:
+        """Get or create Together client with validation."""
+        try:
+            # If api_key from node is empty, try to get from .env
+            final_api_key = api_key.strip() if api_key and api_key.strip() else os.getenv('TOGETHER_API_KEY')
+
+            if not final_api_key:
+                logger.error("No API key provided in node or .env file")
+                raise ValueError("No API key found. Please provide an API key in the node or set TOGETHER_API_KEY in .env file")
+
+            if self.client is None:
+                # Set the API key in environment
+                os.environ["TOGETHER_API_KEY"] = final_api_key
+                self.client = Together()
+            return self.client
+        except Exception as e:
+            logger.error(f"Failed to initialize Together client: {str(e)}")
+            return None
+
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -69,13 +88,6 @@ def INPUT_TYPES(cls):
     CATEGORY = "image"
     OUTPUT_NODE = True
 
-    def get_api_key(self, provided_key: str) -> str:
-        """Get API key with validation."""
-        api_key = provided_key or os.getenv('TOGETHER_API_KEY')
-        if not api_key:
-            raise ValueError("API key not provided. Please provide an API key or set TOGETHER_API_KEY environment variable.")
-        return api_key
-
     def b64_to_image(self, b64_string: str) -> torch.Tensor:
         """Convert base64 string to torch tensor image."""
         try:
@@ -141,13 +153,13 @@ def generate_image(self, prompt: str, api_key: str, width: int, height: int,
             if num_images < 1 or num_images > 4:
                 raise ValueError(f"Number of images must be between 1 and 4. Got {num_images}")
 
-            # Initialize API client
-            api_key = self.get_api_key(api_key)
-            if self.client is None:
-                self.client = Together(api_key=api_key)
+            # Get client with validation and .env fallback
+            client = self.get_client(api_key)
+            if client is None:
+                return None, "Error: Failed to initialize API client. Please check your API key."
 
             # Make API call
-            response = self.client.images.generate(
+            response = client.images.generate(
                 prompt=prompt,
                 model="black-forest-labs/FLUX.1-schnell-Free",
                 width=width,

together_vision_node.py

@@ -27,6 +27,41 @@ def __init__(self):
         self.client = None
         self.last_request_time = 0
         self.min_request_interval = 1.0  # Minimum seconds between requests
+        self.cache = {}  # Response cache
+
+    def validate_api_key(self, api_key: str) -> bool:
+        """Validate if API key is present and well-formed."""
+        if not api_key or len(api_key.strip()) == 0:
+            logger.error("API key is missing or empty")
+            return False
+        return True
+
+    def get_client(self, api_key: str) -> Optional[Together]:
+        """Get or create Together client with validation."""
+        try:
+            # If api_key from node is empty, try to get from .env
+            final_api_key = api_key.strip() if api_key and api_key.strip() else os.getenv('TOGETHER_API_KEY')
+
+            if not final_api_key:
+                logger.error("No API key provided in node or .env file")
+                raise ValueError("No API key found. Please provide an API key in the node or set TOGETHER_API_KEY in .env file")
+
+            if self.client is None:
+                # Set the API key in environment
+                os.environ["TOGETHER_API_KEY"] = final_api_key
+                self.client = Together()
+            return self.client
+        except Exception as e:
+            logger.error(f"Failed to initialize Together client: {str(e)}")
+            return None
+
+    def get_cached_response(self, cache_key: str) -> Optional[str]:
+        """Get cached response if available."""
+        return self.cache.get(cache_key)
+
+    def cache_response(self, cache_key: str, response: str):
+        """Cache the API response."""
+        self.cache[cache_key] = response
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -177,10 +212,10 @@ def process_image(self, model_name: str, api_key: str, system_prompt: str, user_
             }
             actual_model = model_mapping[model_name]
 
-            # Initialize API client
-            api_key = self.get_api_key(api_key)
-            if self.client is None:
-                self.client = Together(api_key=api_key)
+            # Get client with validation and .env fallback
+            client = self.get_client(api_key)
+            if client is None:
+                return ("Error: Failed to initialize API client. Please check your API key.",)
 
             # Prepare messages
             messages = [{"role": "system", "content": system_prompt}]
@@ -207,7 +242,7 @@ def process_image(self, model_name: str, api_key: str, system_prompt: str, user_
 
             try:
                 # API call with timeout handling
-                response = self.client.chat.completions.create(
+                response = client.chat.completions.create(
                     model=actual_model,
                     messages=messages,
                     temperature=temperature,