fix comfyui (#683)

PaddlePaddle · Aug 20, 2024 · 5af011f · 5af011f
1 parent 3bea70f
commit 5af011f
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/comfyui/ComfyUI_ppdiffusers/basic_nodes.py b/comfyui/ComfyUI_ppdiffusers/basic_nodes.py
@@ -36,7 +36,11 @@ def __init__(self):
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {"images": ("IMAGE",), "filename_prefix": ("STRING", {"default": "ComfyUI"})},
+            "required": {
+                "images": ("IMAGE",), 
+                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                "censor": ("BOOLEAN", {"default": True})
+            },
             "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
         }
 
@@ -59,15 +63,18 @@ def censor_image(self, image):
         print(response)
         return response["result"]["pass"]
 
-    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+    def save_images(self, images, censor=True, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
             filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
         )
         results = list()
         for (batch_number, image) in enumerate(images):
             img = Image.fromarray(image)
-            pass_censor = self.censor_image(img)
+            if censor:
+                pass_censor = self.censor_image(img)
+            else:
+                pass_censor = True
             # breakpoint()
             if pass_censor:
                 metadata = None

diff --git a/comfyui/ComfyUI_ppdiffusers/sd_pipe_nodes.py b/comfyui/ComfyUI_ppdiffusers/sd_pipe_nodes.py
@@ -40,7 +40,6 @@ def INPUT_TYPES(cls):
     def load_checkpoint(self, ckpt_name):
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
         pipe = StableDiffusionPipeline.from_single_file(ckpt_path)
-        pipe = pipe.to(dtype=paddle.float16)
         return (pipe,)
 
 

diff --git a/comfyui/ComfyUI_ppdiffusers/sdxl_pipe_nodes.py b/comfyui/ComfyUI_ppdiffusers/sdxl_pipe_nodes.py
@@ -41,7 +41,6 @@ def INPUT_TYPES(cls):
     def load_checkpoint(self, ckpt_name):
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
         pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path)
-        pipe = pipe.to(dtype=paddle.float16)
         return (pipe,)