Add Recraft Text to Vector node, add Save SVG node to handle its output (#53)

comfy_api_nodes/apis/recraft_api.py

@@ -11,11 +11,14 @@ from pydantic import BaseModel, Field, conint
 class RecraftStyle:
     def __init__(self, style: str, substyle: str=None):
         self.style = style
+        if substyle == "None":
+            substyle = None
         self.substyle = substyle
 
 
 class RecraftIO:
     STYLEV3 = "RECRAFT_V3_STYLE"
+    SVG = "SVG" # TODO: if acceptable, move into ComfyUI's typing class
 
 
 class RecraftStyleV3(str, Enum):

comfy_api_nodes/nodes_recraft.py

@@ -19,8 +19,76 @@ from comfy_api_nodes.nodes_api import (
     bytesio_to_image_tensor,
     download_url_to_bytesio,
 )
+import folder_paths
 
+import os
 import torch
+from io import BytesIO
+
+
+class SVG:
+    """
+    Stores SVG representations via a list of BytesIO objects.
+    """
+    def __init__(self, data: list[BytesIO]):
+        self.data = data
+
+
+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 = ()
+    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()
+        for batch_number, svg_bytes in enumerate(svg.data):
+            # NOTE: no way to do metadata for SVG right now, maybe figure this out later
+            # metadata = None
+            # if not args.disable_metadata:
+            #     metadata = PngInfo()
+            #     if prompt is not None:
+            #         metadata.add_text("prompt", json.dumps(prompt))
+            #     if extra_pnginfo is not None:
+            #         for x in extra_pnginfo:
+            #             metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            file = f"{filename_with_batch_num}_{counter:05}_.svg"
+            with open(os.path.join(full_output_folder, file), 'wb') as svg_file:
+                svg_bytes.seek(0)
+                svg_file.write(svg_bytes.read())
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            counter += 1
+        return (None,)
 
 
 class RecraftStyleV3RealisticImageNode:
@@ -197,21 +265,132 @@ class RecraftTextToImageNode:
         return (output_image,)
 
 
+class RecraftTextToVectorNode:
+    """
+    Generates SVG synchronously based on prompt and resolution.
+    """
+
+    RETURN_TYPES = (RecraftIO.SVG,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/image/Recraft"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the image generation.",
+                    },
+                ),
+                "substyle": (get_v3_substyles(RecraftStyleV3.vector_illustration),),
+                "size": (
+                    [res.value for res in RecraftImageSize],
+                    {
+                        "default": RecraftImageSize.res_1024x1024,
+                        "tooltip": "The size of the generated image.",
+                    },
+                ),
+                "n": (
+                    IO.INT,
+                    {
+                        "default": 1,
+                        "min": 1,
+                        "max": 6,
+                        "tooltip": "The number of images to generate.",
+                    },
+                ),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "Seed to determine if node should re-run; actual results are nondeterministic regardless of seed.",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    IO.STRING,
+                    {
+                        "default": "",
+                        "forceInput": True,
+                        "tooltip": "An optional text description of undesired elements on an image.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+            },
+        }
+
+    def api_call(
+        self,
+        prompt: str,
+        substyle: str,
+        size: str,
+        n: int,
+        seed,
+        negative_prompt: str = None,
+        auth_token=None,
+        **kwargs,
+    ):
+        # create RecraftStyle so strings will be formatted properly (i.e. "None" will become None)
+        recraft_style = RecraftStyle(RecraftStyleV3.vector_illustration, substyle=substyle)
+
+        if not negative_prompt:
+            negative_prompt = None
+
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/recraft/image_generation",
+                method=HttpMethod.POST,
+                request_model=RecraftImageGenerationRequest,
+                response_model=RecraftImageGenerationResponse,
+            ),
+            request=RecraftImageGenerationRequest(
+                prompt=prompt,
+                negative_prompts=negative_prompt,
+                model=RecraftModel.recraftv3,
+                size=size,
+                n=n,
+                style=recraft_style.style,
+                substyle=recraft_style.substyle,
+            ),
+            auth_token=auth_token,
+        )
+        response: RecraftImageGenerationResponse = operation.execute()
+        svg_data = []
+        for data in response.data:
+            svg_data.append(download_url_to_bytesio(data.url, timeout=1024))
+
+        return (SVG(svg_data),)
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "RecraftTextToImageNode": RecraftTextToImageNode,
+    "RecraftTextToVectorNode": RecraftTextToVectorNode,
     "RecraftStyleV3RealisticImage": RecraftStyleV3RealisticImageNode,
     "RecraftStyleV3DigitalIllustration": RecraftStyleV3DigitalIllustrationNode,
     "RecraftStyleV3LogoRaster": RecraftStyleV3LogoRasterNode,
-    # "RecraftStyleV3VectorIllustration": RecraftStyleV3VectorIllustrationNode,
+    "SaveSVG": SaveSVGNode,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "RecraftTextToImageNode": "Recraft Text to Image",
+    "RecraftTextToVectorNode": "Recraft Text to Vector",
     "RecraftStyleV3RealisticImage": "Recraft Style - Realistic Image",
     "RecraftStyleV3DigitalIllustration": "Recraft Style - Digital Illustration",
-    "RecraftStyleV3VectorIllustration": "Recraft Style - Vector Illustration",
     "RecraftStyleV3LogoRaster": "Recraft Style - Logo Raster",
+    "SaveSVG": "Save SVG",
 }