Fixed masking issue

gpt_image1.py

@@ -6,7 +6,8 @@ import numpy as np
 from PIL import Image
 import io
 
-class GPTImageEditNode:    
+class GPTImageEditNode:
+    
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -55,11 +56,11 @@ class GPTImageEditNode:
     CATEGORY = "image/ai"
     
     def tensor_to_pil(self, tensor):
-        if len(tensor.shape) == 4:
-            tensor = tensor.squeeze(0)
-        array = (tensor.cpu().numpy() * 255).astype(np.uint8)
-        
-        return Image.fromarray(array)
+        if len(tensor.shape) == 3:
+            array = (tensor.cpu().numpy() * 255).astype(np.uint8)
+            return Image.fromarray(array)
+        else:
+            raise ValueError(f"Expected 3D tensor, got {len(tensor.shape)}D")
     
     def pil_to_tensor(self, image):
         array = np.array(image).astype(np.float32) / 255.0
@@ -72,6 +73,7 @@ class GPTImageEditNode:
             mask = mask.squeeze(0)
         
         array = (mask.cpu().numpy() * 255).astype(np.uint8)
+        
         mask_gray = Image.fromarray(array, mode='L')
         mask_rgba = mask_gray.convert("RGBA")
         mask_rgba.putalpha(mask_gray)
@@ -79,12 +81,15 @@ class GPTImageEditNode:
         return mask_rgba
 
     def invert_mask(self, mask_image):
+        """Invert mask image pixel values"""
         array = np.array(mask_image)
         inverted_array = 255 - array
         return Image.fromarray(inverted_array, mode=mask_image.mode)
-
+    
     def pil_to_bytes(self, image, format="png"):
+        """Convert PIL Image to bytes buffer"""
         buffer = io.BytesIO()
+        
         if format.lower() == 'jpeg':
             image = image.convert('RGB')
         
@@ -95,20 +100,15 @@ class GPTImageEditNode:
     
     def edit_image(self, image, prompt, api_key, background, quality, size, 
                    output_format, output_compression, n_images, mask=None):
-
+        
         try:
-            pil_image = self.tensor_to_pil(image)
-            image_buffer = self.pil_to_bytes(pil_image, "png")
-            
-            if mask is not None:
-                prompt = prompt + " edit in the masked area"
-            
             # Prepare the request
             url = "https://api.openai.com/v1/images/edits"
             headers = {
                 "Authorization": f"Bearer {api_key}"
             }
             
+            # Prepare form data
             files = []
             data = {
                 "model": "gpt-image-1",
@@ -121,12 +121,24 @@ class GPTImageEditNode:
                 "size": size
             }
             
-            files.append(('image[]', (
-                'input.png', 
-                image_buffer, 
-                'image/png'
-            )))
+            # Handle batched images - send all images as reference
+            batch_size = image.shape[0]
+            print(f"Processing {batch_size} images as reference")
             
+            for i in range(batch_size):
+                # Get single image from batch
+                single_image = image[i]
+                pil_image = self.tensor_to_pil(single_image)
+                image_buffer = self.pil_to_bytes(pil_image, "png")
+                
+                # Add to form data
+                files.append(('image[]', (
+                    f'input_{i}.png', 
+                    image_buffer, 
+                    'image/png'
+                )))
+            
+            # Add mask if provided
             if mask is not None:
                 pil_mask = self.mask_to_pil(mask)
                 inverted_mask = self.invert_mask(pil_mask)
@@ -137,36 +149,46 @@ class GPTImageEditNode:
                     'image/png'
                 )))
             
+            # Make the API request
             response = requests.post(url, headers=headers, data=data, files=files)
             
             # Check response
             if response.status_code != 200:
                 error_msg = f"API Error {response.status_code}: {response.text}"
                 print(f"GPT Image Edit Error: {error_msg}")
-                return (image,)
+                # Return first image from batch on error
+                return (image[0:1],)
             
+            # Parse response
             result = response.json()
             
+            # Process the first generated image
             if result['data']:
                 b64_json = result['data'][0]['b64_json']
                 image_bytes = base64.b64decode(b64_json)
                 
+                # Convert to PIL Image
                 pil_image = Image.open(io.BytesIO(image_bytes))
+                
+                # Convert back to tensor
                 output_tensor = self.pil_to_tensor(pil_image)
                 
                 return (output_tensor,)
             else:
                 print("No images returned from API")
-                return (image,)
+                return (image[0:1],)
                 
         except Exception as e:
             print(f"GPT Image Edit Error: {str(e)}")
-            return (image,)
+            # Return first image from batch on error
+            return (image[0:1],)
     
     @classmethod
     def IS_CHANGED(cls, **kwargs):
+        # Always re-execute when prompt changes
         return kwargs.get("prompt", "")
 
+# Node registration
 NODE_CLASS_MAPPINGS = {
     "GPTImageEditNode": GPTImageEditNode
 }