Update More Director Tools

nai.py

@@ -56,8 +56,61 @@ class NovelAIAPI:
         else:
             raise Exception(f"Image generation failed: {response.status_code}, {response.text}")
 
-    def toSketchArt(self,image):
-        ...
+    def toSketchArt(self,input_image):
+        input_image = input_image.squeeze(0)
+        input_image = (input_image * 255).byte().numpy()
+        input_image = Image.fromarray(input_image)
+        buffer = io.BytesIO()
+        input_image.save(buffer, format="PNG")
+        image_payload = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        data = {
+            "height": input_image.height,
+            "width": input_image.width,
+            "req_type": "sketch",
+            "image": image_payload
+        }
+        data_json = json.dumps(data)
+        response = requests.post(self.base_url + "/augment-image", headers=self.headers, data=data_json)
+        if response.status_code == 200:
+            # Handle successful response
+            with zipfile.ZipFile(io.BytesIO(response.content)) as zip_ref:
+                with zip_ref.open("image_0.png") as image_file:
+                    image_data = image_file.read()
+                    # print(image_data)
+                    print("Generated image lineart")
+                    # image = Image.open(io.BytesIO(image_data))
+                    # tensor_image = torch.unsqueeze(torch.tensor(image).permute(2, 0, 1), 0).float() / 255
+                    return (self.convert_image_to_tensor(image_data),)
+        else:
+            raise Exception(f"Image generation failed: {response.status_code}, {response.text}")
+    def declutter(self,input_image):
+        input_image = input_image.squeeze(0)
+        input_image = (input_image * 255).byte().numpy()
+        input_image = Image.fromarray(input_image)
+        buffer = io.BytesIO()
+        input_image.save(buffer, format="PNG")
+        image_payload = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        data = {
+            "height": input_image.height,
+            "width": input_image.width,
+            "req_type": "declutter",
+            "image": image_payload
+        }
+        data_json = json.dumps(data)
+        response = requests.post(self.base_url + "/augment-image", headers=self.headers, data=data_json)
+        if response.status_code == 200:
+            # Handle successful response
+            with zipfile.ZipFile(io.BytesIO(response.content)) as zip_ref:
+                with zip_ref.open("image_0.png") as image_file:
+                    image_data = image_file.read()
+                    # print(image_data)
+                    #print("Generated image lineart")
+                    # image = Image.open(io.BytesIO(image_data))
+                    # tensor_image = torch.unsqueeze(torch.tensor(image).permute(2, 0, 1), 0).float() / 255
+                    return (self.convert_image_to_tensor(image_data),)
+        else:
+            raise Exception(f"Image generation failed: {response.status_code}, {response.text}")
+
     def generate_image(self, input_text, model, width, height, scale, sampler, steps, n_samples=1, ucPreset=0,
                        qualityToggle=True, sm=False, sm_dyn=False, dynamic_thresholding=False, controlnet_strength=1, legacy=False, add_original_image=True,
                        cfg_rescale=0, noise_schedule="native", legacy_v3_extend=False,input_image=None, seed=0, negative_prompt="",

nodes.py

@@ -136,16 +136,52 @@ class NovelAILineart:
         return {
                 "required": {
                     "input_image": ("IMAGE",),
-
                 },
             }
 
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "tolineart"
 
-
     def tolineart(self, input_image=None):
         if (NAI_API_KEY == "<KEY>"):
             raise Exception("API key not set,please configure your API key in config.py in the plugin directory.")
         api = NovelAIAPI(api_key=NAI_API_KEY)
-        return (api.toLineArt(input_image))
+        return (api.toLineArt(input_image))
+class NovelAISketch:
+    CATEGORY = "NovelAI_Sketch_Processor"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+                "required": {
+                    "input_image": ("IMAGE",),
+                },
+            }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "tosketch"
+
+    def tosketch(self, input_image=None):
+        if (NAI_API_KEY == "<KEY>"):
+            raise Exception("API key not set,please configure your API key in config.py in the plugin directory.")
+        api = NovelAIAPI(api_key=NAI_API_KEY)
+        return (api.toSketchArt(input_image))
+class NovelAIDeclutter:
+    CATEGORY = "NovelAI_Declutter_Processor"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+                "required": {
+                    "input_image": ("IMAGE",),
+                },
+            }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "declutter"
+
+    def declutter(self, input_image=None):
+        if (NAI_API_KEY == "<KEY>"):
+            raise Exception("API key not set,please configure your API key in config.py in the plugin directory.")
+        api = NovelAIAPI(api_key=NAI_API_KEY)
+        return (api.declutter(input_image))

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_nai_api"
 description = "A node that can use Nai in Comfyui"
-version = "1.1.0"
+version = "1.1.4"
 license = {file = "LICENSE"}
 
 [project.urls]