[SD] Add fix for using latest diffusers + add scribble variant to Stencil (#1191)

* [SD] Add Scribble variant in Stencil

-- This commit adds scribble variant in Stencil.

Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>

* [SD] Use latest diffusers

-- This commit points back to the latest diffusers and updates the
   processing script to tackle the Pix2Pix import issue.

Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>

---------

Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>
Co-authored-by: Abhishek Varma <abhishek@nod-labs.com>

apps/stable_diffusion/src/utils/stable_args.py

@@ -297,7 +297,7 @@ p.add_argument(
 
 p.add_argument(
     "--use_stencil",
-    choices=["canny", "openpose"],
+    choices=["canny", "openpose", "scribble"],
     help="Enable the stencil feature.",
 )
 

apps/stable_diffusion/src/utils/stencils/stencil_utils.py

@@ -114,6 +114,9 @@ def controlnet_hint_conversion(
         case "openpose":
             print("Detecting human pose")
             controlnet_hint = hint_openpose(image)
+        case "scribble":
+            print("Working with scribble")
+            controlnet_hint = hint_scribble(image)
         case _:
             return None
     controlnet_hint = controlnet_hint_shaping(
@@ -171,3 +174,13 @@ def hint_openpose(
         detected_map, _ = stencil["openpose"](input_image)
         detected_map = HWC3(detected_map)
         return detected_map
+
+
+# Stencil 3. Scribble.
+def hint_scribble(image: Image.Image):
+    with torch.no_grad():
+        input_image = np.array(image)
+
+        detected_map = np.zeros_like(input_image, dtype=np.uint8)
+        detected_map[np.min(input_image, axis=2) < 127] = 255
+        return detected_map

apps/stable_diffusion/web/ui/img2img_ui.py

@@ -70,7 +70,7 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                         use_stencil = gr.Dropdown(
                             label="Stencil model",
                             value="None",
-                            choices=["None", "canny", "openpose"],
+                            choices=["None", "canny", "openpose", "scribble"],
                         )
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():

process_skipfiles.py

@@ -7,14 +7,35 @@ import fileinput
 from pathlib import Path
 
 # Diffusers 0.13.1 fails with transformers __init.py errros in BLIP. So remove it for now until we fork it
-pix2pix_file = Path(
-    get_python_lib()
-    + "/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py"
+pix2pix_init = Path(get_python_lib() + "/diffusers/__init__.py")
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "Pix2Pix" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
+pix2pix_init = Path(get_python_lib() + "/diffusers/pipelines/__init__.py")
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "Pix2Pix" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
+pix2pix_init = Path(
+    get_python_lib() + "/diffusers/pipelines/stable_diffusion/__init__.py"
 )
-if pix2pix_file.exists():
-    print("Removing..%s", pix2pix_file)
-    pix2pix_file.unlink()
-
+for line in fileinput.input(pix2pix_init, inplace=True):
+    if "StableDiffusionPix2PixZeroPipeline" in line:
+        if not line.startswith("#"):
+            print(f"#{line}", end="")
+        else:
+            print(f"{line[1:]}", end="")
+    else:
+        print(line, end="")
 
 path_to_skipfiles = Path(get_python_lib() + "/torch/_dynamo/skipfiles.py")
 

requirements.txt

@@ -16,7 +16,7 @@ parameterized
 
 # Add transformers, diffusers and scipy since it most commonly used
 transformers
-diffusers @ git+https://github.com/nod-ai/diffusers@stable_stencil_diffusers_v0p14p0
+diffusers
 scipy
 ftfy
 gradio < 3.21.0 # temporary workaround for TypeError: issubclass() arg 1 must be a class