Added Recraft Controls + Recraft Color RGB nodes (#57)

comfy_api_nodes/apis/recraft_api.py

@@ -8,6 +8,68 @@ from typing import Optional
 from pydantic import BaseModel, Field, conint
 
 
+class RecraftColor:
+    def __init__(self, r: int, g: int, b: int):
+        self.color = [r, g, b]
+
+    def create_api_model(self):
+        return RecraftColorObject(rgb=self.color)
+
+
+class RecraftColorChain:
+    def __init__(self):
+        self.colors: list[RecraftColor] = []
+
+    def get_first(self):
+        if len(self.colors) > 0:
+            return self.colors[0]
+        return None
+
+    def add(self, color: RecraftColor):
+        self.colors.append(color)
+
+    def create_api_model(self):
+        if not self.colors:
+            return None
+        colors_api = [x.create_api_model() for x in self.colors]
+        return colors_api
+
+    def clone(self):
+        c = RecraftColorChain()
+        for color in self.colors:
+            c.add(color)
+        return c
+
+    def clone_and_merge(self, other: RecraftColorChain):
+        c = self.clone()
+        for color in other.colors:
+            c.add(color)
+        return c
+
+
+class RecraftControls:
+    def __init__(self, colors: RecraftColorChain=None, background_color: RecraftColorChain=None,
+                 artistic_level: int=None, no_text: bool=None):
+        self.colors = colors
+        self.background_color = background_color
+        self.artistic_level = artistic_level
+        self.no_text = no_text
+
+    def create_api_model(self):
+        if self.colors is None and self.background_color is None and self.artistic_level is None and self.no_text is None:
+            return None
+        colors_api = None
+        background_color_api = None
+        if self.colors:
+            colors_api = self.colors.create_api_model()
+        if self.background_color:
+            first_background = self.background_color.get_first()
+            background_color_api = first_background.create_api_model() if first_background else None
+
+        return RecraftControlsObject(colors=colors_api, background_color=background_color_api,
+                                             artistic_level=self.artistic_level, no_text=self.no_text)
+
+
 class RecraftStyle:
     def __init__(self, style: str, substyle: str=None):
         self.style = style
@@ -19,6 +81,8 @@ class RecraftStyle:
 class RecraftIO:
     STYLEV3 = "RECRAFT_V3_STYLE"
     SVG = "SVG" # TODO: if acceptable, move into ComfyUI's typing class
+    COLOR = "RECRAFT_COLOR"
+    CONTROLS = "RECRAFT_CONTROLS"
 
 
 class RecraftStyleV3(str, Enum):
@@ -160,6 +224,17 @@ class RecraftImageSize(str, Enum):
     res_1707x1024 = '1707x1024'
 
 
+class RecraftColorObject(BaseModel):
+    rgb: list[int] = Field(..., description='An array of 3 integer values in range of 0...255 defining RGB Color Model')
+
+
+class RecraftControlsObject(BaseModel):
+    colors: Optional[list[RecraftColorObject]] = Field(None, description='An array of preferable colors')
+    background_color: Optional[RecraftColorObject] = Field(None, description='Use given color as a desired background color')
+    no_text: Optional[bool] = Field(None, description='Do not embed text layouts')
+    artistic_level: Optional[conint(ge=0, le=5)] = Field(None, description='Defines artistic tone of your image. At a simple level, the person looks straight at the camera in a static and clean style. Dynamic and eccentric levels introduce movement and creativity. The value should be in range [0..5].')
+
+
 class RecraftImageGenerationRequest(BaseModel):
     prompt: str = Field(..., description='The text prompt describing the image to generate')
     size: RecraftImageSize = Field(..., description='The size of the generated image (e.g., "1024x1024")')
@@ -168,8 +243,8 @@ class RecraftImageGenerationRequest(BaseModel):
     model: Optional[RecraftModel] = Field(RecraftModel.recraftv3, description='The model to use for generation (e.g., "recraftv3")')
     style: Optional[str] = Field(None, description='The style to apply to the generated image (e.g., "digital_illustration")')
     substyle: Optional[str] = Field(None, description='The substyle to apply to the generated image, depending on the style input')
+    controls: Optional[RecraftControlsObject] = Field(None, description='A set of custom parameters to tweak generation process')
     # text_layout
-    # controls
 
 
 class RecraftReturnedObject(BaseModel):

comfy_api_nodes/nodes_recraft.py

@@ -7,6 +7,9 @@ from comfy_api_nodes.apis.recraft_api import (
     RecraftModel,
     RecraftStyle,
     RecraftStyleV3,
+    RecraftColor,
+    RecraftColorChain,
+    RecraftControls,
     RecraftIO,
     get_v3_substyles,
 )
@@ -91,6 +94,75 @@ class SaveSVGNode:
         return (None,)
 
 
+class RecraftColorRGBNode:
+    """
+    Create Recraft Color by choosing specific RGB values.
+    """
+
+    RETURN_TYPES = (RecraftIO.COLOR,)
+    RETURN_NAMES = ("recraft_color",)
+    FUNCTION = "create_color"
+    CATEGORY = "api node/image/Recraft"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "r": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 255,
+                    "tooltip": "Red value of color."
+                }),
+                "g": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 255,
+                    "tooltip": "Green value of color."
+                }),
+                "b": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 255,
+                    "tooltip": "Blue value of color."
+                }),
+            },
+            "optional": {
+                "recraft_color": (RecraftIO.COLOR,),
+            }
+        }
+
+    def create_color(self, r: int, g: int, b: int, recraft_color: RecraftColorChain=None):
+        recraft_color = recraft_color.clone() if recraft_color else RecraftColorChain()
+        recraft_color.add(RecraftColor(r, g, b))
+        return (recraft_color, )
+
+
+class RecraftControlsNode:
+    """
+    Create Recraft Controls for customizing Recraft generation.
+    """
+
+    RETURN_TYPES = (RecraftIO.CONTROLS,)
+    RETURN_NAMES = ("recraft_controls",)
+    FUNCTION = "create_controls"
+    CATEGORY = "api node/image/Recraft"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+            },
+            "optional": {
+                "colors": (RecraftIO.COLOR,),
+                "background_color": (RecraftIO.COLOR,),
+            }
+        }
+
+    def create_controls(self, colors: RecraftColorChain=None, background_color: RecraftColorChain=None):
+        return (RecraftControls(colors=colors, background_color=background_color), )
+
+
 class RecraftStyleV3RealisticImageNode:
     """
     Select realistic_image style and optional substyle.
@@ -209,6 +281,12 @@ class RecraftTextToImageNode:
                         "tooltip": "An optional text description of undesired elements on an image.",
                     },
                 ),
+                "recraft_controls": (
+                    RecraftIO.CONTROLS,
+                    {
+                        "tooltip": "Optional additional controls over the generation via the Recraft Controls node."
+                    },
+                ),
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
@@ -223,6 +301,7 @@ class RecraftTextToImageNode:
         seed,
         recraft_style: RecraftStyle = None,
         negative_prompt: str = None,
+        recraft_controls: RecraftControls = None,
         auth_token=None,
         **kwargs,
     ):
@@ -230,6 +309,10 @@ class RecraftTextToImageNode:
         if recraft_style is None:
             recraft_style = default_style
 
+        controls_api = None
+        if recraft_controls:
+            controls_api = recraft_controls.create_api_model()
+
         if not negative_prompt:
             negative_prompt = None
 
@@ -248,6 +331,7 @@ class RecraftTextToImageNode:
                 n=n,
                 style=recraft_style.style,
                 substyle=recraft_style.substyle,
+                controls=controls_api,
             ),
             auth_token=auth_token,
         )
@@ -325,6 +409,12 @@ class RecraftTextToVectorNode:
                         "tooltip": "An optional text description of undesired elements on an image.",
                     },
                 ),
+                "recraft_controls": (
+                    RecraftIO.CONTROLS,
+                    {
+                        "tooltip": "Optional additional controls over the generation via the Recraft Controls node."
+                    },
+                ),
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
@@ -339,12 +429,17 @@ class RecraftTextToVectorNode:
         n: int,
         seed,
         negative_prompt: str = None,
+        recraft_controls: RecraftControls = 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)
 
+        controls_api = None
+        if recraft_controls:
+            controls_api = recraft_controls.create_api_model()
+
         if not negative_prompt:
             negative_prompt = None
 
@@ -363,6 +458,7 @@ class RecraftTextToVectorNode:
                 n=n,
                 style=recraft_style.style,
                 substyle=recraft_style.substyle,
+                controls=controls_api,
             ),
             auth_token=auth_token,
         )
@@ -382,6 +478,8 @@ NODE_CLASS_MAPPINGS = {
     "RecraftStyleV3RealisticImage": RecraftStyleV3RealisticImageNode,
     "RecraftStyleV3DigitalIllustration": RecraftStyleV3DigitalIllustrationNode,
     "RecraftStyleV3LogoRaster": RecraftStyleV3LogoRasterNode,
+    "RecraftColorRGB": RecraftColorRGBNode,
+    "RecraftControls": RecraftControlsNode,
     "SaveSVG": SaveSVGNode,
 }
 
@@ -392,5 +490,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "RecraftStyleV3RealisticImage": "Recraft Style - Realistic Image",
     "RecraftStyleV3DigitalIllustration": "Recraft Style - Digital Illustration",
     "RecraftStyleV3LogoRaster": "Recraft Style - Logo Raster",
+    "RecraftColorRGB": "Recraft Color RGB",
+    "RecraftControls": "Recraft Controls",
     "SaveSVG": "Save SVG",
 }