diff --git a/nai.py b/nai.py index 043ffb8..d88ac33 100644 --- a/nai.py +++ b/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="", diff --git a/nodes.py b/nodes.py index c111619..07556c6 100644 --- a/nodes.py +++ b/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 == ""): 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)) \ No newline at end of file + 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 == ""): + 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 == ""): + 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)) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 50ee9ae..aaa3079 100644 --- a/pyproject.toml +++ b/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]