Merge pull request #4 from JosephThomasParker/feature/img2img

Add img2img node

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

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)
 

__init__.py

@@ -1,3 +1,3 @@
 from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes.py

@@ -28,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}),
             }
         }
@@ -37,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"
 
@@ -75,7 +98,7 @@ class 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. * image_tensor.cpu().numpy()
+    i = 255.0 * image_tensor.cpu().numpy()
     img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
 
     # Save the image to a BytesIO object (in memory) rather than to a file
@@ -83,41 +106,41 @@ def image_to_base64(image_tensor):
     img.save(buffered, format="PNG")
 
     # Encode the image as base64
-    encoded_string = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
     return encoded_string
 
+
 def resize_for_inpainting(pixels, mask=None):
-    print(type(pixels))
     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")
+    # 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]
+        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):
+    # 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]
+    Get image size from a size image, i.e. assumed input size is [H, W, C]
     """
-    print(type(pixels))
-    print(np.shape(pixels))
     x = (pixels.shape[0] // 64) * 64
     y = (pixels.shape[1] // 64) * 64
     return x, y
 
+
 class DrawThingsImg2Img:
     def __init__(self):
         pass
@@ -132,11 +155,37 @@ class DrawThingsImg2Img:
                 "model": ("STRING", {"default": "flux_1_dev_q8p.ckpt"}),
                 "prompt": ("STRING", {"default": ""}),
                 "seed": ("INT", {"default": 42}),
-                "width": ("INT", {"default": 512}),
-                "height": ("INT", {"default": 512}),
-                "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"}),
+                "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},
+                ),
             }
         }
 
@@ -144,7 +193,9 @@ class DrawThingsImg2Img:
     RETURN_NAMES = ("generated_image",)
     FUNCTION = "generate_image"
 
-    def generate_image(self, images, model, prompt, seed, width, height, guidance_scale, sampler, steps):
+    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"
 
@@ -155,7 +206,6 @@ class DrawThingsImg2Img:
 
         height, width = get_image_size(images_resized[0])
 
-
         payload = {
             "model": model,
             "prompt": prompt,
@@ -166,40 +216,16 @@ class DrawThingsImg2Img:
             "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()
-#        print("Dia duit!")
-#        #print(data)
-#        print(type(data))
-#        print(type(payload))
-
-        # Path to your PNG image file
-        #image_path = "/Users/jparker/data/sd_outputs/.people/marbro/inputs/1_512sq.JPG"
-
-        # Read the image and encode it as base64
-        #with open(image_path, "rb") as image_file:
-        #    encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
-
-        #print(payload)
-
         response = requests.post(api_url, json=payload)
 
-        data = response.json()
-        print(data)
         # Raise an error if the request failed
         response.raise_for_status()
 
         # Parse the JSON response
         data = response.json()
-        print(data)
 
         # Process the images (assuming they are base64 encoded or raw binary data)
         images = []
@@ -215,11 +241,11 @@ class DrawThingsImg2Img:
 
 
 NODE_CLASS_MAPPINGS = {
-  "DrawThingsTxt2Img": DrawThingsTxt2Img,
-  "DrawThingsImg2Img": DrawThingsImg2Img,
+    "DrawThingsTxt2Img": DrawThingsTxt2Img,
+    "DrawThingsImg2Img": DrawThingsImg2Img,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-  "DrawThingsTxt2Img": "Draw Things Txt2Img", 
-  "DrawThingsImg2Img": "Draw Things Img2Img",
+    "DrawThingsTxt2Img": "Draw Things Txt2Img",
+    "DrawThingsImg2Img": "Draw Things Img2Img",
 }