Upload 49 files

README.md

@@ -11,6 +11,8 @@ license: mit
 duplicated_from:
 - multimodalart/flux-lora-the-explorer
 - gokaygokay/FLUX-Prompt-Generator
+- jiuface/FLUX.1-dev-Controlnet-Union
+- DamarJati/FLUX.1-DEV-Canny
 models: 
   - black-forest-labs/FLUX.1-dev
   - alvdansen/frosting_lane_flux

app.py

@@ -12,7 +12,8 @@ import time
 
 from mod import (models, clear_cache, get_repo_safetensors, change_base_model,
                  description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras,
-                 get_trigger_word, pipe, enhance_prompt)
+                 get_trigger_word, enhance_prompt, pipe, controlnet, num_cns, set_control_union_image,
+                 get_control_union_mode, set_control_union_mode, get_control_params)
 from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json,
                   download_my_lora, get_all_lora_tupled_list, apply_lora_prompt,
                   update_loras)
@@ -64,26 +65,42 @@ def update_selection(evt: gr.SelectData, width, height):
     )
 
 @spaces.GPU(duration=70)
-def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
+def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
     progress(0, desc="Start Inference.")
     with calculateDuration("Generating image"):
         # Generate image
-        image = pipe(
-            prompt=prompt_mash,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
+        modes, images, scales = get_control_params()
+        if not cn_on or controlnet is None or len(modes) == 0:
+            image = pipe(
+                prompt=prompt_mash,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                joint_attention_kwargs={"scale": lora_scale},
+            ).images[0]
+        else:
+            
+            image = pipe(
+                prompt=prompt_mash,
+                control_image=images,
+                control_mode=modes,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                controlnet_conditioning_scale=scales,
+                generator=generator,
+                joint_attention_kwargs={"scale": lora_scale},
+            ).images[0]
     return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
-              lora_scale, lora_json, progress=gr.Progress(track_tqdm=True)):
+              lora_scale, lora_json, cn_on, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None and not is_valid_lora(lora_json):
         gr.Info("LoRA isn't selected.")
     #    raise gr.Error("You must select a LoRA before proceeding.")
@@ -123,7 +140,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     
     progress(1, desc="Preparing Inference.")
 
-    image = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    image = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress)
     if is_valid_lora(lora_json):
         pipe.unfuse_lora()
         pipe.unload_lora_weights()
@@ -318,7 +335,24 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                                     lora_download = [None] * num_loras
                                     for i in range(num_loras):
                                         lora_download[i] = gr.Button(f"Get and set LoRA to {int(i+1)}")
-     
+
+                    with gr.Accordion("ControlNet", open=True):
+                        with gr.Column():
+                            cn_on = gr.Checkbox(False, label="Use ControlNet")
+                            cn_mode = [None] * num_cns
+                            cn_scale = [None] * num_cns
+                            cn_image = [None] * num_cns
+                            cn_res = [None] * num_cns
+                            cn_num = [None] * num_cns
+                            for i in range(num_cns):
+                                with gr.Group():
+                                    with gr.Row():
+                                        cn_mode[i] = gr.Dropdown(label=f"ControlNet {int(i+1)} Mode", choices=get_control_union_mode(), value=get_control_union_mode()[0], allow_custom_value=False)
+                                        cn_scale[i] = gr.Slider(label=f"ControlNet {int(i+1)} Weight", minimum=0.0, maximum=1.0, step=0.01, value=0.75)
+                                        cn_res[i] = gr.Slider(label=f"ControlNet {int(i+1)} Preprocess resolution", minimum=128, maximum=512, value=384, step=1)
+                                        cn_num[i] = gr.Number(i, visible=False)
+                                    cn_image[i] = gr.Image(type="pil", label="Control Image", height=256)
+
     gallery.select(
         update_selection,
         inputs=[width, height],
@@ -336,16 +370,21 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=change_base_model,
-        inputs=[model_name],
+        inputs=[model_name, cn_on],
         outputs=[result]
     ).success(
         fn=run_lora,
         inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
-                 lora_scale, lora_repo_json], 
+                 lora_scale, lora_repo_json, cn_on], 
         outputs=[result, seed]
     )
 
-    model_name.change(change_base_model, [model_name], [result])
+    gr.on(
+        triggers=[model_name.change, cn_on.change],
+        fn=change_base_model,
+        inputs=[model_name, cn_on],
+        outputs=[result]
+    )
     prompt_enhance.click(enhance_prompt, [prompt], [prompt])
 
     gr.on(
@@ -382,6 +421,16 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
         ).success(apply_lora_prompt, [lora_info[i]], [lora_trigger[i]], queue=False, show_api=False
         ).success(compose_lora_json, [lora_repo_json, lora_num[i], lora_repo[i], lora_wt[i], lora_weights[i], lora_trigger[i]], [lora_repo_json], queue=False, show_api=False)
 
+    for i, m in enumerate(cn_mode):
+        gr.on(
+            triggers=[cn_mode[i].change, cn_scale[i].change],
+            fn=set_control_union_mode,
+            inputs=[cn_num[i], cn_mode[i], cn_scale[i]],
+            outputs=[cn_on],
+            queue=True,
+            show_api=False,
+        )
+        cn_image[i].upload(set_control_union_image, [cn_num[i], cn_mode[i], cn_image[i], height, width, cn_res[i]], [cn_image[i]])
 
     tagger_generate_from_image.click(
             lambda: ("", "", ""), None, [v2_series, v2_character, prompt], queue=False, show_api=False,

cv_utils.py

@@ -0,0 +1,18 @@
+import cv2
+import numpy as np
+
+MAX_IMAGE_SIZE = 512
+
+def resize_image(input_image, resolution=MAX_IMAGE_SIZE, interpolation=None):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / max(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    if interpolation is None:
+        interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
+    img = cv2.resize(input_image, (W, H), interpolation=interpolation)
+    return img

depth_estimator.py

@@ -0,0 +1,14 @@
+import numpy as np
+import PIL.Image
+from controlnet_aux.util import HWC3
+from transformers import pipeline
+
+from cv_utils import resize_image
+
+
+class DepthEstimator:
+    def __init__(self):
+        self.model = pipeline("depth-estimation")
+
+    def __call__(self, image: np.ndarray, **kwargs) -> PIL.Image.Image:
+        return image

image_datasets/canny_dataset.py

@@ -0,0 +1,59 @@
+import os
+import pandas as pd
+import numpy as np
+from PIL import Image
+import torch
+from torch.utils.data import Dataset, DataLoader
+import json
+import random
+import cv2
+
+
+def canny_processor(image, low_threshold=100, high_threshold=200):
+    image = np.array(image)
+    image = cv2.Canny(image, low_threshold, high_threshold)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    canny_image = Image.fromarray(image)
+    return canny_image
+
+
+def c_crop(image):
+    width, height = image.size
+    new_size = min(width, height)
+    left = (width - new_size) / 2
+    top = (height - new_size) / 2
+    right = (width + new_size) / 2
+    bottom = (height + new_size) / 2
+    return image.crop((left, top, right, bottom))
+
+class CustomImageDataset(Dataset):
+    def __init__(self, img_dir, img_size=512):
+        self.images = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if '.jpg' in i or '.png' in i]
+        self.images.sort()
+        self.img_size = img_size
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        try:
+            img = Image.open(self.images[idx])
+            img = c_crop(img)
+            img = img.resize((self.img_size, self.img_size))
+            hint = canny_processor(img)
+            img = torch.from_numpy((np.array(img) / 127.5) - 1)
+            img = img.permute(2, 0, 1)
+            hint = torch.from_numpy((np.array(hint) / 127.5) - 1)
+            hint = hint.permute(2, 0, 1)
+            json_path = self.images[idx].split('.')[0] + '.json'
+            prompt = json.load(open(json_path))['caption']
+            return img, hint, prompt
+        except Exception as e:
+            print(e)
+            return self.__getitem__(random.randint(0, len(self.images) - 1))
+
+
+def loader(train_batch_size, num_workers, **args):
+    dataset = CustomImageDataset(**args)
+    return DataLoader(dataset, batch_size=train_batch_size, num_workers=num_workers)

image_datasets/dataset.py

@@ -0,0 +1,45 @@
+import os
+import pandas as pd
+import numpy as np
+from PIL import Image
+import torch
+from torch.utils.data import Dataset, DataLoader
+import json
+import random
+
+def c_crop(image):
+    width, height = image.size
+    new_size = min(width, height)
+    left = (width - new_size) / 2
+    top = (height - new_size) / 2
+    right = (width + new_size) / 2
+    bottom = (height + new_size) / 2
+    return image.crop((left, top, right, bottom))
+
+class CustomImageDataset(Dataset):
+    def __init__(self, img_dir, img_size=512):
+        self.images = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if '.jpg' in i or '.png' in i]
+        self.images.sort()
+        self.img_size = img_size
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        try:
+            img = Image.open(self.images[idx])
+            img = c_crop(img)
+            img = img.resize((self.img_size, self.img_size))
+            img = torch.from_numpy((np.array(img) / 127.5) - 1)
+            img = img.permute(2, 0, 1)
+            json_path = self.images[idx].split('.')[0] + '.json'
+            prompt = json.load(open(json_path))['caption']
+            return img, prompt
+        except Exception as e:
+            print(e)
+            return self.__getitem__(random.randint(0, len(self.images) - 1))
+
+
+def loader(train_batch_size, num_workers, **args):
+    dataset = CustomImageDataset(**args)
+    return DataLoader(dataset, batch_size=train_batch_size, num_workers=num_workers, shuffle=True)

image_segmentor.py

@@ -0,0 +1,34 @@
+import cv2
+import numpy as np
+import PIL.Image
+import torch
+from controlnet_aux.util import HWC3, ade_palette
+from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
+
+from cv_utils import resize_image
+
+
+class ImageSegmentor:
+    
+    def __init__(self):
+        self.image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
+        self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
+
+    @torch.no_grad()
+    def __call__(self, image: np.ndarray, **kwargs) -> PIL.Image.Image:
+        detect_resolution = kwargs.pop("detect_resolution", 512)
+        image_resolution = kwargs.pop("image_resolution", 512)
+        image = HWC3(image)
+        image = resize_image(image, resolution=detect_resolution)
+        image = PIL.Image.fromarray(image)
+
+        pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
+        outputs = self.image_segmentor(pixel_values)
+        seg = self.image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
+        for label, color in enumerate(ade_palette()):
+            color_seg[seg == label, :] = color
+        color_seg = color_seg.astype(np.uint8)
+
+        color_seg = resize_image(color_seg, resolution=image_resolution, interpolation=cv2.INTER_NEAREST)
+        return PIL.Image.fromarray(color_seg)

mod.py

@@ -2,9 +2,12 @@ import gradio as gr
 import torch
 import spaces
 from diffusers import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline
+from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel
 from pathlib import Path
 import gc
 import subprocess
+from PIL import Image
 
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -31,7 +34,17 @@ models = [
 
 
 num_loras = 3
-
+num_cns = 2
+# Initialize the base model
+base_model = models[0]
+controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+controlnet = None
+control_images = [None] * num_cns
+control_modes = [-1] * num_cns
+control_scales = [0] * num_cns
+last_model = models[0]
+last_cn_on = False
 
 def is_repo_name(s):
     import re
@@ -70,26 +83,169 @@ def get_repo_safetensors(repo_id: str):
     else: return gr.update(value=files[0], choices=files)
 
 
-# Initialize the base model
-base_model = models[0]
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-last_model = models[0]
-
-def change_base_model(repo_id: str, progress=gr.Progress(track_tqdm=True)):
+# https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny
+# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
+# https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
+def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm=True)):
     global pipe
+    global controlnet
     global last_model
+    global last_cn_on
     try:
-        if repo_id == last_model or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return
-        progress(0, desc=f"Loading model: {repo_id}")
-        clear_cache()
-        pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-        last_model = repo_id
-        progress(1, desc=f"Model loaded: {repo_id}")
+        if (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return
+        if cn_on:
+            progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
+            print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
+            clear_cache()
+            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=torch.bfloat16)
+            controlnet = FluxMultiControlNetModel([controlnet_union])
+            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=torch.bfloat16)
+            last_model = repo_id
+            progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
+            print(f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
+        else:
+            progress(0, desc=f"Loading model: {repo_id}")
+            print(f"Loading model: {repo_id}")
+            clear_cache()
+            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+            last_model = repo_id
+            progress(1, desc=f"Model loaded: {repo_id}")
+            print(f"Model loaded: {repo_id}")
     except Exception as e:
         print(e)
     return gr.update(visible=True)
 
 
+# https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny/blob/main/app.py
+def resize_image(image, target_width, target_height, crop=True):
+    from image_datasets.canny_dataset import c_crop
+    if crop:
+        image = c_crop(image)  # Crop the image to square
+        original_width, original_height = image.size
+
+        # Resize to match the target size without stretching
+        scale = max(target_width / original_width, target_height / original_height)
+        resized_width = int(scale * original_width)
+        resized_height = int(scale * original_height)
+
+        image = image.resize((resized_width, resized_height), Image.LANCZOS)
+        
+        # Center crop to match the target dimensions
+        left = (resized_width - target_width) // 2
+        top = (resized_height - target_height) // 2
+        image = image.crop((left, top, left + target_width, top + target_height))
+    else:
+        image = image.resize((target_width, target_height), Image.LANCZOS)
+    
+    return image
+
+
+# https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union/blob/main/app.py
+controlnet_union_modes = {
+    "None": -1,
+    #"scribble_hed": 0,
+    "canny": 0, # supported
+    "mlsd": 0, #supported
+    "tile": 1, #supported
+    "depth_midas": 2, # supported
+    "blur": 3, # supported
+    "openpose": 4,  # supported
+    "gray": 5,  # supported
+    "low_quality": 6,  # supported
+}
+
+
+def get_control_params():
+    modes = []
+    images = []
+    scales = []
+    for i, mode in enumerate(control_modes):
+        if mode == -1 or control_images[i] is None: continue
+        modes.append(control_modes[i])
+        images.append(control_images[i])
+        scales.append(control_scales[i])
+    return modes, images, scales
+
+
+from preprocessor import Preprocessor
+def preprocess_image(image: Image.Image, control_mode: str, height: int, width: int, preprocess_resolution: int):
+    image_resolution = max(width, height)
+    image_before = resize_image(image, image_resolution, image_resolution, True)
+    # generated control_
+    print("start to generate control image")
+    preprocessor = Preprocessor()
+    if control_mode == "depth_midas":
+        preprocessor.load("Midas")
+        control_image = preprocessor(
+            image=image_before,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+    if control_mode == "openpose":
+        preprocessor.load("Openpose")
+        control_image = preprocessor(
+            image=image_before,
+            hand_and_face=True,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+    if control_mode == "canny":
+        preprocessor.load("Canny")
+        control_image = preprocessor(
+            image=image_before,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+
+    if control_mode == "mlsd":
+        preprocessor.load("MLSD")
+        control_image = preprocessor(
+            image=image_before,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+
+    if control_mode == "scribble_hed":
+        preprocessor.load("HED")
+        control_image = preprocessor(
+            image=image_before,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+    
+    if control_mode == "low_quality" or  control_mode == "gray" or  control_mode == "blur" or control_mode == "tile":
+        control_image = image_before
+        image_width = 768
+        image_height = 768
+    else:
+        # make sure control image size is same as resized_image
+        image_width, image_height =  control_image.size
+    
+    image_after = resize_image(control_image, width, height, True)
+    print(f"generate control image success: {image_width}x{image_height} => {width}x{height}")
+
+    return image_after
+
+
+def get_control_union_mode():
+    return list(controlnet_union_modes.keys())
+
+
+def set_control_union_mode(i: int, mode: str, scale: str):
+    global control_modes
+    global control_scales
+    control_modes[i] = controlnet_union_modes.get(mode, 0)
+    control_scales[i] = scale
+    if mode != "None": return True
+    else: return gr.update(visible=True)
+
+
+def set_control_union_image(i: int, mode: str, image: Image.Image, height: int, width: int, preprocess_resolution: int):
+    global control_images
+    control_images[i] = preprocess_image(image, mode, height, width, preprocess_resolution)
+    return control_images[i]
+
+
 def compose_lora_json(lorajson: list[dict], i: int, name: str, scale: float, filename: str, trigger: str):
     lorajson[i]["name"] = str(name) if name != "None" else ""
     lorajson[i]["scale"] = float(scale)
@@ -112,6 +268,7 @@ def get_trigger_word(lorajson: list[dict]):
             trigger += ", " + d["trigger"]
     return trigger
 
+
 # https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora
 # https://github.com/huggingface/diffusers/issues/4919
 def fuse_loras(pipe, lorajson: list[dict]):
@@ -139,13 +296,12 @@ def fuse_loras(pipe, lorajson: list[dict]):
     #pipe.unload_lora_weights()
 
 
-
-
-
 def description_ui():
     gr.Markdown(
         """
 - Mod of [multimodalart/flux-lora-the-explorer](https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer),
+ [jiuface/FLUX.1-dev-Controlnet-Union](https://huggingface.co/spaces/jiuface/),
+ [DamarJati/FLUX.1-DEV-Canny](https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny),
  [gokaygokay/FLUX-Prompt-Generator](https://huggingface.co/spaces/gokaygokay/FLUX-Prompt-Generator).
 """
     )

preprocessor.py

@@ -0,0 +1,84 @@
+import gc
+
+import numpy as np
+import PIL.Image
+import torch
+import torchvision
+from controlnet_aux import (
+    CannyDetector,
+    ContentShuffleDetector,
+    HEDdetector,
+    LineartAnimeDetector,
+    LineartDetector,
+    MidasDetector,
+    MLSDdetector,
+    NormalBaeDetector,
+    OpenposeDetector,
+    PidiNetDetector,
+)
+from controlnet_aux.util import HWC3
+
+from cv_utils import resize_image
+from depth_estimator import DepthEstimator
+from image_segmentor import ImageSegmentor
+
+from kornia.core import Tensor
+
+# load preprocessor
+
+# HED = HEDdetector.from_pretrained("lllyasviel/Annotators")
+Midas = MidasDetector.from_pretrained("lllyasviel/Annotators")
+MLSD = MLSDdetector.from_pretrained("lllyasviel/Annotators")
+Canny = CannyDetector()
+OPENPOSE =  OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+
+    def __init__(self):
+        self.model = None
+        self.name = ""
+
+    def load(self, name: str) -> None:
+        if name == self.name:
+            return
+        
+        if name == "Midas":
+            self.model = Midas
+        elif name == "MLSD":
+            self.model =MLSD
+        elif name == "Openpose":
+            self.model = OPENPOSE
+        elif name == "Canny":
+            self.model = Canny
+        else:
+            raise ValueError
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        if self.name == "Canny" or self.name == "MLSD":
+            detect_resolution = kwargs.pop("detect_resolution")
+            image_resolution = kwargs.pop("image_resolution", 512)
+            image = np.array(image)
+            image = HWC3(image)
+            image = resize_image(image, resolution=detect_resolution)
+            image = self.model(image, **kwargs)
+            image = np.array(image)
+            image = HWC3(image)
+            image = resize_image(image, resolution=image_resolution)
+            return PIL.Image.fromarray(image).convert('RGB')
+
+        else:
+            detect_resolution = kwargs.pop("detect_resolution", 512)
+            image_resolution = kwargs.pop("image_resolution", 512)
+            image = np.array(image)
+            image = HWC3(image)
+            image = resize_image(image, resolution=detect_resolution)
+            image = self.model(image, **kwargs)
+            image = np.array(image)
+            image = HWC3(image)
+            image = resize_image(image, resolution=image_resolution)
+            return PIL.Image.fromarray(image)

requirements.txt

@@ -8,5 +8,7 @@ transformers
 peft
 sentencepiece
 timm
-xformers
-einops
+einops
+controlnet-aux
+kornia
+numpy