Move ControlNet up so pipe available for concept population

app.py

@@ -44,8 +44,9 @@ my_token = os.environ['api_key']
 
 pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", revision="fp16", torch_dtype=torch.float16, use_auth_token=my_token).to("cuda")
 
+
 def check_prompt(prompt):
-    SPAM_WORDS = ['Д', 'boob', 'boobs', 'breast', 'breasts'] # phasing this out 
+    SPAM_WORDS = ['Д', 'oob', 'reast'] # only necessary to limit spam
     for spam_word in SPAM_WORDS:
         if spam_word in prompt:
             return False
@@ -81,6 +82,77 @@ def load_learned_embed_in_clip(learned_embeds_path, text_encoder, tokenizer, tok
   return token
 
 
+
+# ----- ControlNet Canny Edges -----------------------------------------------------------------
+
+import gradio as gr
+from PIL import Image
+import numpy as np
+import cv2
+
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import UniPCMultistepScheduler
+import torch
+
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
+controlnet_pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+)
+
+controlnet_pipe.scheduler = UniPCMultistepScheduler.from_config(controlnet_pipe.scheduler.config)
+controlnet_pipe.enable_model_cpu_offload()
+controlnet_pipe.enable_xformers_memory_efficient_attention()
+
+def controlnet_edges(canny_input_prompt, input_image, input_low_threshold, input_high_threshold, input_invert):
+    np_image = np.array(input_image)
+
+    output_image = input_image
+    numpy_image = np.array(output_image)
+
+    low_threshold = 80
+    high_threshold = 100
+    canny_1 = cv2.Canny(numpy_image, input_low_threshold, input_high_threshold)
+    canny_1 = canny_1[:, :, None]
+    canny_1 = np.concatenate([canny_1, canny_1, canny_1], axis=2)
+    if input_invert:
+      canny_1 = 255 - canny_1
+
+    canny_2 = Image.fromarray(canny_1)
+
+    prompt = canny_input_prompt
+    generator = torch.Generator(device="cpu").manual_seed(2)
+
+    output_image = controlnet_pipe(
+        prompt,
+        canny_2,
+        negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
+        generator=generator,
+        num_inference_steps=20,
+    )
+
+    return output_image[0][0]
+
+
+canny_input_prompt = gr.inputs.Textbox(label="Enter a single word or phrase")
+canny_input_image = gr.inputs.Image()
+canny_input_low_threshold = gr.inputs.Slider(minimum=0, maximum=1000, step=1, label="Lower Threshold:", default=100)
+canny_input_high_threshold = gr.inputs.Slider(minimum=0, maximum=1000, step=1, label="Upper Threshold:", default=200)
+canny_input_invert = gr.inputs.Checkbox(label="Invert Image")
+canny_outputs = gr.outputs.Image(type="pil")
+
+# make and launch the gradio app...
+controlnet_canny_interface = gr.Interface(fn=controlnet_edges, inputs=[canny_input_prompt, canny_input_image, canny_input_low_threshold, canny_input_high_threshold, canny_input_invert], outputs=canny_outputs, title='Canny Edge Tracing', allow_flagging='never')
+# controlnet_canny_interface.launch()
+
+
+
+
+
+
+
+# ----- Load All models / concepts -----------------------------------------------------------------
+
+
 ahx_model_list = [model for model in models_list if "ahx" in model.modelId]
 ahx_dropdown_list = [model for model in models_list if "ahx-model" in model.modelId]
 
@@ -119,6 +191,7 @@ for model in ahx_model_list:
   #if token cannot be loaded, skip it
   try:
     learned_token = load_learned_embed_in_clip(f"{model_id}/learned_embeds.bin", pipe.text_encoder, pipe.tokenizer, token_name)
+    # _learned_token_controlnet = load_learned_embed_in_clip(f"{model_id}/learned_embeds.bin", controlnet_pipe.text_encoder, controlnet_pipe.tokenizer, token_name)
   except: 
     continue
   model_content["token"] = learned_token
@@ -412,69 +485,6 @@ canny_interface = gr.Interface(fn=canny_process_image, inputs=[canny_input_image
 
 
 
-# ----- ControlNet Canny Edges -----------------------------------------------------------------
-
-# import gradio as gr
-# from PIL import Image
-# import numpy as np
-# import cv2
-
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-from diffusers import UniPCMultistepScheduler
-import torch
-
-controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-controlnet_pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
-)
-
-controlnet_pipe.scheduler = UniPCMultistepScheduler.from_config(controlnet_pipe.scheduler.config)
-controlnet_pipe.enable_model_cpu_offload()
-controlnet_pipe.enable_xformers_memory_efficient_attention()
-
-def controlnet_edges(canny_input_prompt, input_image, input_low_threshold, input_high_threshold, input_invert):
-    np_image = np.array(input_image)
-
-    output_image = input_image
-    numpy_image = np.array(output_image)
-
-    low_threshold = 80
-    high_threshold = 100
-    canny_1 = cv2.Canny(numpy_image, input_low_threshold, input_high_threshold)
-    canny_1 = canny_1[:, :, None]
-    canny_1 = np.concatenate([canny_1, canny_1, canny_1], axis=2)
-    if input_invert:
-      canny_1 = 255 - canny_1
-
-    canny_2 = Image.fromarray(canny_1)
-
-    prompt = canny_input_prompt
-    generator = torch.Generator(device="cpu").manual_seed(2)
-
-    output_image = controlnet_pipe(
-        prompt,
-        canny_2,
-        negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
-        generator=generator,
-        num_inference_steps=20,
-    )
-
-    return output_image[0][0]
-
-
-canny_input_prompt = gr.inputs.Textbox(label="Enter a single word or phrase")
-canny_input_image = gr.inputs.Image()
-canny_input_low_threshold = gr.inputs.Slider(minimum=0, maximum=1000, step=1, label="Lower Threshold:", default=100)
-canny_input_high_threshold = gr.inputs.Slider(minimum=0, maximum=1000, step=1, label="Upper Threshold:", default=200)
-canny_input_invert = gr.inputs.Checkbox(label="Invert Image")
-canny_outputs = gr.outputs.Image(type="pil")
-
-# make and launch the gradio app...
-controlnet_canny_interface = gr.Interface(fn=controlnet_edges, inputs=[canny_input_prompt, canny_input_image, canny_input_low_threshold, canny_input_high_threshold, canny_input_invert], outputs=canny_outputs, title='Canny Edge Tracing', allow_flagging='never')
-# controlnet_canny_interface.launch()
-
-
-
 
 # ----- Launch Tabs -----------------------------------------------------------------