From ced625d9b016e974f9fdaebb5d9348dedaa8475a Mon Sep 17 00:00:00 2001 From: Slava Cherkasov Date: Wed, 20 Aug 2025 10:19:32 +0200 Subject: [PATCH] smart image save now supports GIF --- nodes.py | 248 ++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 245 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index 02e0547..710574c 100644 --- a/nodes.py +++ b/nodes.py @@ -1607,6 +1607,7 @@ class SmartSaveAnimatedPNG: # Keep lossy_quality for potential future use, but set default high "lossless": ("BOOLEAN", {"default": True}), "save_metadata": ("BOOLEAN", {"default": True}), + "format": (["APNG", "PNGif", "GIF"], {"default": "APNG"}), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -1615,15 +1616,256 @@ class SmartSaveAnimatedPNG: FUNCTION = "save_apng" OUTPUT_NODE = True CATEGORY = "SmartImageTools" + + def save_animated_gif(self, images, filename, fps, loop, optimize, quality, preserve_transparency, alpha_threshold): + try: + from PIL import Image + import numpy as np + import torch + except ImportError as e: + raise RuntimeError(f"Required libraries not available: {e}. Please install Pillow.") - def save_apng(self, images, fps, filename_prefix="ComfyUI", lossless=True, save_metadata=True, prompt=None, extra_pnginfo=None): + # Ensure filename has .gif extension + if not filename.lower().endswith('.gif'): + if filename.endswith('.'): + filename = filename + 'gif' + else: + filename = filename + '.gif' + + # Convert relative path to absolute path within ComfyUI's output directory + output_dir = folder_paths.get_output_directory() + full_path = os.path.join(output_dir, filename) + + # Create directories if they don't exist + directory = os.path.dirname(full_path) + if directory: + os.makedirs(directory, exist_ok=True) + + # Convert images to PIL format + pil_frames = [] + has_transparent_pixel = False + used_colors = set() + alpha_thresh = alpha_threshold * 255 + + if isinstance(images, torch.Tensor) and images.dim() == 4: + # Handle batch of images + batch_size = images.shape[0] + else: + # If not a batch tensor, treat as list + batch_size = len(images) + + for i in range(batch_size): + if isinstance(images, torch.Tensor): + img = images[i] + else: + img = images[i] + + if isinstance(img, torch.Tensor): + img = img.detach().cpu() + if img.dim() == 3: + if img.shape[0] in [1, 3, 4]: # CHW + img = img.permute(1, 2, 0) + if img.max() <= 1.0: + img = img * 255 + img = img.round().clamp(0, 255).byte().numpy() + mode = 'RGBA' if img.shape[-1] == 4 else 'RGB' if img.shape[-1] == 3 else 'L' if img.shape[-1] == 1 else None + if mode is None: + raise ValueError(f"Unsupported tensor channel count: {img.shape[-1]}") + if mode == 'L': + img = np.repeat(img[:, :, np.newaxis], 3, axis=-1) + mode = 'RGB' + pil_img = Image.fromarray(img, mode) + elif isinstance(img, np.ndarray): + if img.max() <= 1.0: + img = img * 255 + img = np.round(img).clip(0, 255).astype(np.uint8) + mode = 'RGBA' if img.shape[-1] == 4 else 'RGB' if img.shape[-1] == 3 else 'L' if img.shape[-1] == 1 else None + if mode is None: + raise ValueError(f"Unsupported array channel count: {img.shape[-1]}") + if mode == 'L': + img = np.repeat(img[:, :, np.newaxis], 3, axis=-1) + mode = 'RGB' + pil_img = Image.fromarray(img, mode) + elif isinstance(img, Image.Image): + pil_img = img + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + # Convert to RGBA for consistency if transparency might be present + if pil_img.mode not in ['RGB', 'RGBA']: + pil_img = pil_img.convert('RGB') + + if preserve_transparency and pil_img.mode == 'RGBA': + data = np.array(pil_img) + alpha = data[:, :, 3] + if np.any(alpha < alpha_thresh): + has_transparent_pixel = True + opaque_mask = alpha >= alpha_thresh + opaque_rgb = data[opaque_mask, :3] + for color in opaque_rgb: + used_colors.add(tuple(map(int, color))) + + pil_frames.append(pil_img) + + if not pil_frames: + raise ValueError("No valid images provided") + + if not preserve_transparency or not has_transparent_pixel: + # Composite to RGB with white background + for i in range(len(pil_frames)): + if pil_frames[i].mode == 'RGBA': + background = Image.new('RGB', pil_frames[i].size, (255, 255, 255)) + background.paste(pil_frames[i], mask=pil_frames[i].getchannel('A')) + pil_frames[i] = background + pil_images = pil_frames + else: + # Find global key_color not used in any opaque pixel + key_color = None + attempts = 0 + while attempts < 10000 and key_color is None: + r = random.randint(0, 255) + g = random.randint(0, 255) + b = random.randint(0, 255) + if (r, g, b) not in used_colors: + key_color = (r, g, b) + attempts += 1 + + if key_color is None: + key_color = (255, 0, 255) + # Replace occurrences of key_color in opaque areas + for i in range(len(pil_frames)): + if pil_frames[i].mode == 'RGBA': + data = np.array(pil_frames[i]) + alpha = data[:, :, 3] + match = (data[:, :, 0] == 255) & (data[:, :, 1] == 0) & (data[:, :, 2] == 255) & (alpha >= alpha_thresh) + data[match, 0] = 254 + pil_frames[i] = Image.fromarray(data, 'RGBA') + + # Create RGB frames with key_color in transparent areas + rgb_frames = [] + for frame in pil_frames: + if frame.mode == 'RGBA': + data = np.array(frame) + transparent_mask = data[:, :, 3] < alpha_thresh + rgb_data = data[:, :, :3] + rgb_data[transparent_mask] = key_color + frame_rgb = Image.fromarray(rgb_data, 'RGB') + else: + frame_rgb = frame.convert('RGB') + rgb_frames.append(frame_rgb) + + # Create large image for global palette + widths, heights = zip(*(f.size for f in rgb_frames)) + max_width = max(widths) + total_height = sum(heights) + 32 # Extra for key block + large_img = Image.new('RGB', (max_width, total_height)) + y = 0 + for f in rgb_frames: + large_img.paste(f, (0, y)) + y += f.height + key_block = Image.new('RGB', (32, 32), key_color) + large_img.paste(key_block, (0, y)) + + # Quantize large image + large_quant = large_img.quantize(colors=256, method=Image.Quantize.MAXCOVERAGE) + global_palette = large_quant.getpalette()[:256*3] + + # Find trans_index + trans_index = -1 + for i in range(256): + if (global_palette[i*3] == key_color[0] and + global_palette[i*3 + 1] == key_color[1] and + global_palette[i*3 + 2] == key_color[2]): + trans_index = i + break + + if trans_index == -1: + # Fallback to per-frame palettes + pil_images = [] + for frame_rgb in rgb_frames: + frame_p = frame_rgb.quantize(colors=256) + palette = frame_p.getpalette()[:256*3] + for j in range(256): + if (palette[j*3] == key_color[0] and + palette[j*3 + 1] == key_color[1] and + palette[j*3 + 2] == key_color[2]): + frame_p.info['transparency'] = j + frame_p.info['disposal'] = 2 + break + pil_images.append(frame_p) + if pil_images: + pil_images[0].info['background'] = pil_images[0].info.get('transparency', 0) + else: + # Use global palette + palette_img = Image.new('P', (1, 1)) + palette_img.putpalette(global_palette) + pil_images = [] + for frame_rgb in rgb_frames: + frame_p = frame_rgb.quantize(palette=palette_img) + frame_p.info['transparency'] = trans_index + frame_p.info['disposal'] = 2 + pil_images.append(frame_p) + if pil_images: + pil_images[0].info['background'] = trans_index + + # Calculate duration per frame in milliseconds + duration = int(1000 / fps) + + # Save as animated GIF + pil_images[0].save( + full_path, + save_all=True, + append_images=pil_images[1:], + duration=duration, + loop=loop, + optimize=optimize, + quality=quality + ) + + return (full_path,) + + def save_apng(self, images, fps, filename_prefix="ComfyUI", lossless=True, save_metadata=True, format="APNG", prompt=None, extra_pnginfo=None): full_output_folder, filename, counter, subfolder, filename_prefix_ = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) # Extract the base filename part from the potentially path-like filename_prefix_ base_filename = os.path.basename(filename_prefix_) # Use a single filename, not numbered sequence - file = f"{base_filename}.apng" # Use the extracted base name + if format == "APNG": + ext = "apng" + elif format == "PNGif": + ext = "pngif" + else: + ext = "gif" + file = f"{base_filename}.{ext}" # Use the extracted base name file_path = os.path.join(full_output_folder, file) + # GIF branch: delegate to save_animated_gif + if format == "GIF": + results = list() + try: + relative_filename = os.path.join(subfolder, file) if subfolder else file + self.save_animated_gif( + images=images, + filename=relative_filename, + fps=fps, + loop=0, + optimize=True, + quality=95, + preserve_transparency=True, + alpha_threshold=0.01, + ) + + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + + except Exception as e: + print(f"Error saving GIF file: {e}") + raise + + return {"ui": {"apng": results}} + num_frames = images.shape[0] fps_int = int(round(fps * 100)) # Multiply by 100 and convert to int @@ -1827,7 +2069,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SmartImagePaletteExtract": "Smart Image Palette Extract", "SmartSemiTransparenceRemove": "Smart Semi-Transparence Remove", "SmartVideoPreviewScaled": "Smart Video Preview Scaled", - "SmartSaveAnimatedPNG": "Smart Save Animated PNG", + "SmartSaveAnimatedPNG": "Smart Save Animated Image", "SmartSavePNG": "Smart Save PNG", "SmartDrawPoints": "Smart Draw Points" } \ No newline at end of file