code · pull · May 9, 2025 · May 9, 2025 · May 9, 2025
diff --git a/comfy/ldm/ace/attention.py b/comfy/ldm/ace/attention.py
@@ -19,6 +19,7 @@
 from torch import nn
 
 import comfy.model_management
+from comfy.ldm.modules.attention import optimized_attention
 
 class Attention(nn.Module):
     def __init__(
@@ -326,10 +327,6 @@ class CustomerAttnProcessor2_0:
     Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
     """
 
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
     def apply_rotary_emb(
         self,
         x: torch.Tensor,
@@ -435,13 +432,9 @@ def __call__(
             attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
 
         # the output of sdp = (batch, num_heads, seq_len, head_dim)
-        # TODO: add support for attn.scale when we move to Torch 2.1
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
-
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = optimized_attention(
+            query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True,
+        ).to(query.dtype)
 
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)

diff --git a/comfy_api_nodes/apis/recraft_api.py b/comfy_api_nodes/apis/recraft_api.py
@@ -81,7 +81,6 @@ def __init__(self, style: str=None, substyle: str=None, style_id: str=None):
 
 class RecraftIO:
     STYLEV3 = "RECRAFT_V3_STYLE"
-    SVG = "SVG" # TODO: if acceptable, move into ComfyUI's typing class
     COLOR = "RECRAFT_COLOR"
     CONTROLS = "RECRAFT_CONTROLS"
 

diff --git a/comfy_api_nodes/nodes_recraft.py b/comfy_api_nodes/nodes_recraft.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 from inspect import cleandoc
 from comfy.utils import ProgressBar
+from comfy_extras.nodes_images import SVG # Added
 from comfy.comfy_types.node_typing import IO
 from comfy_api_nodes.apis.recraft_api import (
     RecraftImageGenerationRequest,
@@ -28,9 +29,6 @@
     resize_mask_to_image,
     validate_string,
 )
-import folder_paths
-import json
-import os
 import torch
 from io import BytesIO
 from PIL import UnidentifiedImageError
@@ -162,102 +160,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             raise Exception("Received output data was not an image; likely an SVG. If you used style_id, make sure it is not a Vector art style.")
 
 
-class SVG:
-    """
-    Stores SVG representations via a list of BytesIO objects.
-    """
-    def __init__(self, data: list[BytesIO]):
-        self.data = data
-
-    def combine(self, other: SVG):
-        return SVG(self.data + other.data)
-
-    @staticmethod
-    def combine_all(svgs: list[SVG]):
-        all_svgs = []
-        for svg in svgs:
-            all_svgs.extend(svg.data)
-        return SVG(all_svgs)
-
-
-class SaveSVGNode:
-    """
-    Save SVG files on disk.
-    """
-
-    def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-
-    RETURN_TYPES = ()
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "save_svg"
-    CATEGORY = "api node/image/Recraft"
-    OUTPUT_NODE = True
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "svg": (RecraftIO.SVG,),
-                "filename_prefix": ("STRING", {"default": "svg/ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
-            },
-            "hidden": {
-                "prompt": "PROMPT",
-                "extra_pnginfo": "EXTRA_PNGINFO"
-            }
-        }
-
-    def save_svg(self, svg: SVG, filename_prefix="svg/ComfyUI", prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
-        results = list()
-
-        # Prepare metadata JSON
-        metadata_dict = {}
-        if prompt is not None:
-            metadata_dict["prompt"] = prompt
-        if extra_pnginfo is not None:
-            metadata_dict.update(extra_pnginfo)
-
-        # Convert metadata to JSON string
-        metadata_json = json.dumps(metadata_dict, indent=2) if metadata_dict else None
-
-        for batch_number, svg_bytes in enumerate(svg.data):
-            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
-            file = f"{filename_with_batch_num}_{counter:05}_.svg"
-
-            # Read SVG content
-            svg_bytes.seek(0)
-            svg_content = svg_bytes.read().decode('utf-8')
-
-            # Inject metadata if available
-            if metadata_json:
-                # Create metadata element with CDATA section
-                metadata_element = f"""  <metadata>
-    <![CDATA[
-{metadata_json}
-    ]]>
-  </metadata>
-"""
-                # Insert metadata after opening svg tag using regex
-                import re
-                svg_content = re.sub(r'(<svg[^>]*>)', r'\1\n' + metadata_element, svg_content)
-
-            # Write the modified SVG to file
-            with open(os.path.join(full_output_folder, file), 'wb') as svg_file:
-                svg_file.write(svg_content.encode('utf-8'))
-
-            results.append({
-                "filename": file,
-                "subfolder": subfolder,
-                "type": self.type
-            })
-            counter += 1
-        return { "ui": { "images": results } }
-
-
 class RecraftColorRGBNode:
     """
     Create Recraft Color by choosing specific RGB values.
@@ -796,8 +698,8 @@ class RecraftTextToVectorNode:
     Generates SVG synchronously based on prompt and resolution.
     """
 
-    RETURN_TYPES = (RecraftIO.SVG,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    RETURN_TYPES = ("SVG",) # Changed
+    DESCRIPTION = cleandoc(__doc__ or "") if 'cleandoc' in globals() else __doc__ # Keep cleandoc if other nodes use it
     FUNCTION = "api_call"
     API_NODE = True
     CATEGORY = "api node/image/Recraft"
@@ -918,8 +820,8 @@ class RecraftVectorizeImageNode:
     Generates SVG synchronously from an input image.
     """
 
-    RETURN_TYPES = (RecraftIO.SVG,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    RETURN_TYPES = ("SVG",) # Changed
+    DESCRIPTION = cleandoc(__doc__ or "") if 'cleandoc' in globals() else __doc__ # Keep cleandoc if other nodes use it
     FUNCTION = "api_call"
     API_NODE = True
     CATEGORY = "api node/image/Recraft"
@@ -1193,7 +1095,6 @@ class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode):
     "RecraftStyleV3InfiniteStyleLibrary": RecraftStyleInfiniteStyleLibrary,
     "RecraftColorRGB": RecraftColorRGBNode,
     "RecraftControls": RecraftControlsNode,
-    "SaveSVG": SaveSVGNode,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -1213,5 +1114,4 @@ class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode):
     "RecraftStyleV3InfiniteStyleLibrary": "Recraft Style - Infinite Style Library",
     "RecraftColorRGB": "Recraft Color RGB",
     "RecraftControls": "Recraft Controls",
-    "SaveSVG": "Save SVG",
 }
diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
@@ -10,6 +10,9 @@
 import numpy as np
 import json
 import os
+import re
+from io import BytesIO
+from inspect import cleandoc
 
 from comfy.comfy_types import FileLocator
 
@@ -190,10 +193,109 @@ def save_images(self, images, fps, compress_level, filename_prefix="ComfyUI", pr
 
         return { "ui": { "images": results, "animated": (True,)} }
 
+class SVG:
+    """
+    Stores SVG representations via a list of BytesIO objects.
+    """
+    def __init__(self, data: list[BytesIO]):
+        self.data = data
+
+    def combine(self, other: 'SVG') -> 'SVG':
+        return SVG(self.data + other.data)
+
+    @staticmethod
+    def combine_all(svgs: list['SVG']) -> 'SVG':
+        all_svgs_list: list[BytesIO] = []
+        for svg_item in svgs:
+            all_svgs_list.extend(svg_item.data)
+        return SVG(all_svgs_list)
+
+class SaveSVGNode:
+    """
+    Save SVG files on disk.
+    """
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    RETURN_TYPES = ()
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "save_svg"
+    CATEGORY = "image/save" # Changed
+    OUTPUT_NODE = True
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "svg": ("SVG",), # Changed
+                "filename_prefix": ("STRING", {"default": "svg/ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "extra_pnginfo": "EXTRA_PNGINFO"
+            }
+        }
+
+    def save_svg(self, svg: SVG, filename_prefix="svg/ComfyUI", prompt=None, extra_pnginfo=None):
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        results = list()
+
+        # Prepare metadata JSON
+        metadata_dict = {}
+        if prompt is not None:
+            metadata_dict["prompt"] = prompt
+        if extra_pnginfo is not None:
+            metadata_dict.update(extra_pnginfo)
+
+        # Convert metadata to JSON string
+        metadata_json = json.dumps(metadata_dict, indent=2) if metadata_dict else None
+
+        for batch_number, svg_bytes in enumerate(svg.data):
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            file = f"{filename_with_batch_num}_{counter:05}_.svg"
+
+            # Read SVG content
+            svg_bytes.seek(0)
+            svg_content = svg_bytes.read().decode('utf-8')
+
+            # Inject metadata if available
+            if metadata_json:
+                # Create metadata element with CDATA section
+                metadata_element = f"""  <metadata>
+                <![CDATA[
+            {metadata_json}
+                ]]>
+            </metadata>
+            """
+                # Insert metadata after opening svg tag using regex with a replacement function
+                def replacement(match):
+                    # match.group(1) contains the captured <svg> tag
+                    return match.group(1) + '\n' + metadata_element
+
+                # Apply the substitution
+                svg_content = re.sub(r'(<svg[^>]*>)', replacement, svg_content, flags=re.UNICODE)
+
+            # Write the modified SVG to file
+            with open(os.path.join(full_output_folder, file), 'wb') as svg_file:
+                svg_file.write(svg_content.encode('utf-8'))
+
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            counter += 1
+        return { "ui": { "images": results } }
+
 NODE_CLASS_MAPPINGS = {
     "ImageCrop": ImageCrop,
     "RepeatImageBatch": RepeatImageBatch,
     "ImageFromBatch": ImageFromBatch,
     "SaveAnimatedWEBP": SaveAnimatedWEBP,
     "SaveAnimatedPNG": SaveAnimatedPNG,
+    "SaveSVGNode": SaveSVGNode,
 }