diff --git a/gpt_image1.py b/gpt_image1.py index f0c7f2c..d7640c9 100644 --- a/gpt_image1.py +++ b/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 }