diff --git a/CHANGELOG.md b/CHANGELOG.md index 0711da6..643610c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +0.3 + +* Added img2img node + 0.2 * Added textbox to change model (must match downloaded model) and dropdown list to choose sampler diff --git a/README.md b/README.md index 6cd9cbc..b666cf3 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ These nodes provide a wrapper for calling [Draw Things](https://drawthings.ai/) **Wait, why?** The Draw Things app has been optimized for Apple hardware and runs roughly x3 faster than ComfyUI generations. But ComfyUI is a flexible and powerful tools, and has some features - like queuing and face swapping - that haven't been implemented in Draw Things. -This simple node calls a local instance of Draw Things through its API and returns the resulting image to ComfyUI. +These simple nodes for txt2img and img2img call a local instance of Draw Things through its API and return the resulting image to ComfyUI. ![](basic_example.png) diff --git a/__init__.py b/__init__.py index 2e96bd6..39a8c6b 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,3 @@ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/basic_example.png b/basic_example.png index 71ca30c..ea4b22c 100644 Binary files a/basic_example.png and b/basic_example.png differ diff --git a/nodes.py b/nodes.py index 5fb882f..5a70323 100644 --- a/nodes.py +++ b/nodes.py @@ -8,6 +8,7 @@ import numpy as np import requests from PIL import Image import io +from io import BytesIO import torch @@ -27,7 +28,28 @@ class DrawThingsTxt2Img: "width": ("INT", {"default": 512}), "height": ("INT", {"default": 512}), "guidance_scale": ("FLOAT", {"default": 3.5}), - "sampler": (["UniPC","DPM++ 2M Karras","Euler Ancestral", "DPM++ SDE Karras", "PLMS", "DDIM", "LCM", "Euler A Substep", "DPM++ SDE Substep", "TCD", "DPM++ 2M Trailing", "Euler A Trailing", "DPM++ SDE Trailing", "DDIM Trailing", "DPM++ 2M AYS", "Euler A AYS", "DPM++ SDE AYS"], {"default": "Euler A Trailing"}), + "sampler": ( + [ + "UniPC", + "DPM++ 2M Karras", + "Euler Ancestral", + "DPM++ SDE Karras", + "PLMS", + "DDIM", + "LCM", + "Euler A Substep", + "DPM++ SDE Substep", + "TCD", + "DPM++ 2M Trailing", + "Euler A Trailing", + "DPM++ SDE Trailing", + "DDIM Trailing", + "DPM++ 2M AYS", + "Euler A AYS", + "DPM++ SDE AYS", + ], + {"default": "Euler A Trailing"}, + ), "steps": ("INT", {"default": 20}), } } @@ -36,7 +58,9 @@ class DrawThingsTxt2Img: RETURN_NAMES = ("generated_image",) FUNCTION = "generate_image" - def generate_image(self, model, prompt, seed, width, height, guidance_scale, sampler, steps): + def generate_image( + self, model, prompt, seed, width, height, guidance_scale, sampler, steps + ): # Call the Draw Things API api_url = "http://127.0.0.1:7860/sdapi/v1/txt2img" @@ -72,6 +96,156 @@ class DrawThingsTxt2Img: return (torch.stack(images),) -NODE_CLASS_MAPPINGS = {"DrawThingsTxt2Img": DrawThingsTxt2Img} +def image_to_base64(image_tensor): + # Convert the image tensor to a NumPy array and scale it to the range 0-255 + i = 255.0 * image_tensor.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) -NODE_DISPLAY_NAME_MAPPINGS = {"DrawThingsTxt2Img": "Draw Things Txt2Img"} + # Save the image to a BytesIO object (in memory) rather than to a file + buffered = BytesIO() + img.save(buffered, format="PNG") + + # Encode the image as base64 + encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8") + return encoded_string + + +def resize_for_inpainting(pixels, mask=None): + x = (pixels.shape[1] // 64) * 64 + y = (pixels.shape[2] // 64) * 64 + # mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear") + + orig_pixels = pixels + pixels = orig_pixels.clone() + if pixels.shape[1] != x or pixels.shape[2] != y: + x_offset = (pixels.shape[1] % 64) // 2 + y_offset = (pixels.shape[2] % 64) // 2 + pixels = pixels[:, x_offset : x + x_offset, y_offset : y + y_offset, :] + # pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset] + # mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset] + + # m = (1.0 - mask.round()).squeeze(1) + # for i in range(3): + # pixels[:,:,:,i] -= 0.5 + # pixels[:,:,:,i] *= m + # pixels[:,:,:,i] += 0.5 + return pixels + + +def get_image_size(pixels): + """ + Get image size from a size image, i.e. assumed input size is [H, W, C] + """ + x = (pixels.shape[0] // 64) * 64 + y = (pixels.shape[1] // 64) * 64 + return x, y + + +class DrawThingsImg2Img: + def __init__(self): + pass + + CATEGORY = "DrawThingsWrapper" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", {"tooltip": "input image"}), + "model": ("STRING", {"default": "flux_1_dev_q8p.ckpt"}), + "prompt": ("STRING", {"default": ""}), + "seed": ("INT", {"default": 42}), + "guidance_scale": ( + "FLOAT", + {"default": 3.5, "min": 0, "max": 25, "step": 0.1}, + ), + "sampler": ( + [ + "UniPC", + "DPM++ 2M Karras", + "Euler Ancestral", + "DPM++ SDE Karras", + "PLMS", + "DDIM", + "LCM", + "Euler A Substep", + "DPM++ SDE Substep", + "TCD", + "DPM++ 2M Trailing", + "Euler A Trailing", + "DPM++ SDE Trailing", + "DDIM Trailing", + "DPM++ 2M AYS", + "Euler A AYS", + "DPM++ SDE AYS", + ], + {"default": "Euler A Trailing"}, + ), + "steps": ("INT", {"default": 20, "min": 1, "max": 150, "step": 1}), + "denoise": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("generated_image",) + FUNCTION = "generate_image" + + def generate_image( + self, images, model, prompt, seed, guidance_scale, sampler, steps, denoise + ): + # Call the Draw Things API + api_url = "http://127.0.0.1:7860/sdapi/v1/img2img" + + encoded_images = [] + images_resized = resize_for_inpainting(images) + for image_tensor in images_resized: + encoded_images.append(image_to_base64(image_tensor)) + + height, width = get_image_size(images_resized[0]) + + payload = { + "model": model, + "prompt": prompt, + "seed": seed, + "width": width, + "height": height, + "guidance_scale": guidance_scale, + "sampler": sampler, + "steps": steps, + "init_images": encoded_images, + "strength": denoise, + } + + response = requests.post(api_url, json=payload) + + # Raise an error if the request failed + response.raise_for_status() + + # Parse the JSON response + data = response.json() + + # Process the images (assuming they are base64 encoded or raw binary data) + images = [] + for img_data in data["images"]: + image_bytes = base64.b64decode(img_data) + # Convert the image data to a Pillow Image object + image = Image.open(io.BytesIO(image_bytes)) + image_np = np.array(image) + # Convert to float32 tensor and normalize + tensor_image = torch.from_numpy(image_np.astype(np.float32) / 255.0) + images.append(tensor_image) + return (torch.stack(images),) + + +NODE_CLASS_MAPPINGS = { + "DrawThingsTxt2Img": DrawThingsTxt2Img, + "DrawThingsImg2Img": DrawThingsImg2Img, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DrawThingsTxt2Img": "Draw Things Txt2Img", + "DrawThingsImg2Img": "Draw Things Img2Img", +}