first commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+*.pt
+*.pth
+*.st

README.md

@@ -0,0 +1,13 @@
+---
+title: UniControl Demo
+emoji: 📚
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: 3.36.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

_app.py

@@ -0,0 +1,1360 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+import config
+
+import cv2
+import einops
+import gradio as gr
+import numpy as np
+import torch
+import random
+import os
+
+from pytorch_lightning import seed_everything
+from annotator.util import resize_image, HWC3
+from annotator.uniformer_base import UniformerDetector
+from annotator.hed import HEDdetector
+from annotator.canny import CannyDetector
+from annotator.midas import MidasDetector
+from annotator.outpainting import Outpainter
+from annotator.openpose import OpenposeDetector
+from annotator.inpainting import Inpainter
+from annotator.grayscale import GrayscaleConverter
+from annotator.blur import Blurrer
+import cvlib as cv
+
+from utils import create_model, load_state_dict
+from lib.ddim_hacked import DDIMSampler
+
+from safetensors.torch import load_file as stload
+from collections import OrderedDict
+
+apply_uniformer = UniformerDetector()
+apply_midas = MidasDetector()
+apply_canny = CannyDetector()
+apply_hed = HEDdetector()
+model_outpainting = Outpainter()
+apply_openpose = OpenposeDetector()
+model_grayscale = GrayscaleConverter()
+model_blur = Blurrer()
+model_inpainting = Inpainter()
+
+
+def midas(img, res):
+    img = resize_image(HWC3(img), res)
+    results = apply_midas(img)
+    return results
+
+
+def outpainting(img, res, height_top_extended, height_down_extended, width_left_extended, width_right_extended):
+    img = resize_image(HWC3(img), res)
+    result = model_outpainting(img, height_top_extended, height_down_extended, width_left_extended, width_right_extended)
+    return result
+
+
+def grayscale(img, res):
+    img = resize_image(HWC3(img), res)
+    result = model_grayscale(img)
+    return result
+
+
+def blur(img, res, ksize):
+    img = resize_image(HWC3(img), res)
+    result = model_blur(img, ksize)
+    return result
+
+
+def inpainting(img, res, height_top_mask, height_down_mask, width_left_mask, width_right_mask):
+    img = resize_image(HWC3(img), res)
+    result = model_inpainting(img, height_top_mask, height_down_mask, width_left_mask, width_right_mask)
+    return result
+
+model = create_model('./models/cldm_v15_unicontrol.yaml').cpu()
+# model_url = 'https://huggingface.co/Robert001/UniControl-Model/resolve/main/unicontrol_v1.1.ckpt'
+model_url = 'https://huggingface.co/Robert001/UniControl-Model/resolve/main/unicontrol_v1.1.st'
+
+ckpts_path='./'
+# model_path = os.path.join(ckpts_path, "unicontrol_v1.1.ckpt")
+model_path = os.path.join(ckpts_path, "unicontrol_v1.1.st")
+
+if not os.path.exists(model_path):
+    from basicsr.utils.download_util import load_file_from_url
+    load_file_from_url(model_url, model_dir=ckpts_path)
+
+model_dict = OrderedDict(stload(model_path, device='cpu'))
+model.load_state_dict(model_dict, strict=False)
+# model.load_state_dict(load_state_dict(model_path, location='cuda'), strict=False)
+model = model.cuda()
+ddim_sampler = DDIMSampler(model)
+
+task_to_name = {'hed': 'control_hed', 'canny': 'control_canny', 'seg': 'control_seg', 'segbase': 'control_seg',
+                'depth': 'control_depth', 'normal': 'control_normal', 'openpose': 'control_openpose',
+                'bbox': 'control_bbox', 'grayscale': 'control_grayscale', 'outpainting': 'control_outpainting',
+                'hedsketch': 'control_hedsketch', 'inpainting': 'control_inpainting', 'blur': 'control_blur',
+                'grayscale': 'control_grayscale'}
+
+name_to_instruction = {"control_hed": "hed edge to image", "control_canny": "canny edge to image",
+                       "control_seg": "segmentation map to image", "control_depth": "depth map to image",
+                       "control_normal": "normal surface map to image", "control_img": "image editing",
+                       "control_openpose": "human pose skeleton to image", "control_hedsketch": "sketch to image",
+                       "control_bbox": "bounding box to image", "control_outpainting": "image outpainting",
+                       "control_grayscale": "gray image to color image", "control_blur": "deblur image to clean image",
+                       "control_inpainting": "image inpainting"}
+
+
+def process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                  strength, scale, seed, eta, low_threshold, high_threshold, condition_mode):
+    with torch.no_grad():
+        img = resize_image(HWC3(input_image), image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = apply_canny(img, low_threshold, high_threshold)
+            detected_map = HWC3(detected_map)
+        else:
+            detected_map = 255 - img
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+        task = 'canny'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [255 - detected_map] + results
+
+
+def process_hed(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps,
+                guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = apply_hed(resize_image(input_image, detect_resolution))
+            detected_map = HWC3(detected_map)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'hed'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_depth(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps,
+                  guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map, _ = apply_midas(resize_image(input_image, detect_resolution))
+            detected_map = HWC3(detected_map)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+        task = 'depth'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_normal(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            _, detected_map = apply_midas(resize_image(input_image, detect_resolution))
+            detected_map = HWC3(detected_map)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+        task = 'normal'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_pose(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps,
+                 guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map, _ = apply_openpose(resize_image(input_image, detect_resolution))
+            detected_map = HWC3(detected_map)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+        task = 'openpose'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_seg(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps,
+                guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        if condition_mode == True:
+            detected_map = apply_uniformer(resize_image(input_image, detect_resolution))
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+        task = 'seg'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+color_dict = {
+    'background': (0, 0, 100),
+    'person': (255, 0, 0),
+    'bicycle': (0, 255, 0),
+    'car': (0, 0, 255),
+    'motorcycle': (255, 255, 0),
+    'airplane': (255, 0, 255),
+    'bus': (0, 255, 255),
+    'train': (128, 128, 0),
+    'truck': (128, 0, 128),
+    'boat': (0, 128, 128),
+    'traffic light': (128, 128, 128),
+    'fire hydrant': (64, 0, 0),
+    'stop sign': (0, 64, 0),
+    'parking meter': (0, 0, 64),
+    'bench': (64, 64, 0),
+    'bird': (64, 0, 64),
+    'cat': (0, 64, 64),
+    'dog': (192, 192, 192),
+    'horse': (32, 32, 32),
+    'sheep': (96, 96, 96),
+    'cow': (160, 160, 160),
+    'elephant': (224, 224, 224),
+    'bear': (32, 0, 0),
+    'zebra': (0, 32, 0),
+    'giraffe': (0, 0, 32),
+    'backpack': (32, 32, 0),
+    'umbrella': (32, 0, 32),
+    'handbag': (0, 32, 32),
+    'tie': (96, 0, 0),
+    'suitcase': (0, 96, 0),
+    'frisbee': (0, 0, 96),
+    'skis': (96, 96, 0),
+    'snowboard': (96, 0, 96),
+    'sports ball': (0, 96, 96),
+    'kite': (160, 0, 0),
+    'baseball bat': (0, 160, 0),
+    'baseball glove': (0, 0, 160),
+    'skateboard': (160, 160, 0),
+    'surfboard': (160, 0, 160),
+    'tennis racket': (0, 160, 160),
+    'bottle': (224, 0, 0),
+    'wine glass': (0, 224, 0),
+    'cup': (0, 0, 224),
+    'fork': (224, 224, 0),
+    'knife': (224, 0, 224),
+    'spoon': (0, 224, 224),
+    'bowl': (64, 64, 64),
+    'banana': (128, 64, 64),
+    'apple': (64, 128, 64),
+    'sandwich': (64, 64, 128),
+    'orange': (128, 128, 64),
+    'broccoli': (128, 64, 128),
+    'carrot': (64, 128, 128),
+    'hot dog': (192, 64, 64),
+    'pizza': (64, 192, 64),
+    'donut': (64, 64, 192),
+    'cake': (192, 192, 64),
+    'chair': (192, 64, 192),
+    'couch': (64, 192, 192),
+    'potted plant': (96, 32, 32),
+    'bed': (32, 96, 32),
+    'dining table': (32, 32, 96),
+    'toilet': (96, 96, 32),
+    'tv': (96, 32, 96),
+    'laptop': (32, 96, 96),
+    'mouse': (160, 32, 32),
+    'remote': (32, 160, 32),
+    'keyboard': (32, 32, 160),
+    'cell phone': (160, 160, 32),
+    'microwave': (160, 32, 160),
+    'oven': (32, 160, 160),
+    'toaster': (224, 32, 32),
+    'sink': (32, 224, 32),
+    'refrigerator': (32, 32, 224),
+    'book': (224, 224, 32),
+    'clock': (224, 32, 224),
+    'vase': (32, 224, 224),
+    'scissors': (64, 96, 96),
+    'teddy bear': (96, 64, 96),
+    'hair drier': (96, 96, 64),
+    'toothbrush': (160, 96, 96)
+}
+
+
+def process_bbox(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                 strength, scale, seed, eta, confidence, nms_thresh, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        if condition_mode == True:
+            bbox, label, conf = cv.detect_common_objects(input_image, confidence=confidence, nms_thresh=nms_thresh)
+            mask = np.zeros((input_image.shape), np.uint8)
+            if len(bbox) > 0:
+                order_area = np.zeros(len(bbox))
+                #     order_final = np.arange(len(bbox))
+                area_all = 0
+                for idx_mask, box in enumerate(bbox):
+                    x_1, y_1, x_2, y_2 = box
+
+                    x_1 = 0 if x_1 < 0 else x_1
+                    y_1 = 0 if y_1 < 0 else y_1
+                    x_2 = input_image.shape[1] if x_2 < 0 else x_2
+                    y_2 = input_image.shape[0] if y_2 < 0 else y_2
+
+                    area = (x_2 - x_1) * (y_2 - y_1)
+                    order_area[idx_mask] = area
+                    area_all += area
+                ordered_area = np.argsort(-order_area)
+
+                for idx_mask in ordered_area:
+                    box = bbox[idx_mask]
+                    x_1, y_1, x_2, y_2 = box
+                    x_1 = 0 if x_1 < 0 else x_1
+                    y_1 = 0 if y_1 < 0 else y_1
+                    x_2 = input_image.shape[1] if x_2 < 0 else x_2
+                    y_2 = input_image.shape[0] if y_2 < 0 else y_2
+
+                    mask[y_1:y_2, x_1:x_2, :] = color_dict[label[idx_mask]]
+            detected_map = mask
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'bbox'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_outpainting(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                        strength, scale, seed, eta, height_top_extended, height_down_extended, width_left_extended, width_right_extended, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = outpainting(input_image, image_resolution, height_top_extended, height_down_extended, width_left_extended, width_right_extended)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'outpainting'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_sketch(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        if condition_mode == True:
+            detected_map = apply_hed(resize_image(input_image, detect_resolution))
+            detected_map = HWC3(detected_map)
+
+            # sketch the hed image
+            retry = 0
+            cnt = 0
+            while retry == 0:
+                threshold_value = np.random.randint(110, 160)
+                kernel_size = 3
+                alpha = 1.5
+                beta = 50
+                binary_image = cv2.threshold(detected_map, threshold_value, 255, cv2.THRESH_BINARY)[1]
+                inverted_image = cv2.bitwise_not(binary_image)
+                smoothed_image = cv2.GaussianBlur(inverted_image, (kernel_size, kernel_size), 0)
+                sketch_image = cv2.convertScaleAbs(smoothed_image, alpha=alpha, beta=beta)
+                if np.sum(sketch_image < 5) > 0.005 * sketch_image.shape[0] * sketch_image.shape[1] or cnt == 5:
+                    retry = 1
+                else:
+                    cnt += 1
+            detected_map = sketch_image
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'hedsketch'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_colorization(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                         strength, scale, seed, eta, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = grayscale(input_image, image_resolution)
+            detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+            detected_map = detected_map[:, :, np.newaxis]
+            detected_map = detected_map.repeat(3, axis=2)
+        else:
+            detected_map = img
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'grayscale'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_deblur(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, ksize, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = blur(input_image, image_resolution, ksize)
+        else:
+            detected_map = img
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'blur'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+def process_inpainting(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                       strength, scale, seed, eta, h_ratio_t, h_ratio_d, w_ratio_l, w_ratio_r, condition_mode):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+        if condition_mode == True:
+            detected_map = inpainting(input_image, image_resolution, h_ratio_t, h_ratio_d, w_ratio_l, w_ratio_r)
+        else:
+            detected_map = img
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        task = 'inpainting'
+        task_dic = {}
+        task_dic['name'] = task_to_name[task]
+        task_instruction = name_to_instruction[task_dic['name']]
+        task_dic['feature'] = model.get_learned_conditioning(task_instruction)[:, :1, :]
+
+        cond = {"c_concat": [control],
+                "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)],
+                "task": task_dic}
+        un_cond = {"c_concat": [control * 0] if guess_mode else [control],
+                   "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
+                    [strength] * 13)
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0,
+                                                                                                           255).astype(
+            np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+############################################################################################################
+
+
+demo = gr.Blocks()
+with demo:
+    #gr.Markdown("UniControl Stable Diffusion Demo")
+    gr.HTML(
+        """
+        <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
+        <h1 style="font-weight: 900; font-size: 2rem; margin: 0rem">
+            UniControl Stable Diffusion Demo
+        </h1>
+        <p style="font-size: 1rem; margin: 0rem">
+           Can Qin <sup>1,2</sup>, Shu Zhang<sup>1</sup>, Ning Yu <sup>1</sup>, Yihao Feng<sup>1</sup>, Xinyi Yang<sup>1</sup>, Yingbo Zhou <sup>1</sup>, Huan Wang <sup>1</sup>, Juan Carlos Niebles<sup>1</sup>, Caiming Xiong <sup>1</sup>, Silvio Savarese <sup>1</sup>, Stefano Ermon <sup>3</sup>, Yun Fu <sup>2</sup>,  Ran Xu <sup>1</sup> 
+        </p>
+        <p style="font-size: 0.8rem; margin: 0rem; line-height: 1em">
+            <sup>1</sup> Salesforce AI <sup>2</sup> Northeastern University  <sup>3</sup> Stanford University
+        </p>
+        <p style="font-size: 0.8rem; margin: 0rem; line-height: 1em">
+        Work done when Can Qin was an intern at Salesforce AI Research.
+        </p>
+        <p style="font-size: 0.9rem; margin: 0rem; line-height: 1.2em; margin-top:1em">
+           <b> ONE compact model for ALL the visual-condition-to-image generation! </b> 
+            <b><a href="https://github.com/salesforce/UniControl">[Github]</a></b> 
+            <b><a href="https://canqin001.github.io/UniControl-Page/">[Website]</a></b> 
+             <b><a href="https://arxiv.org/abs/2305.11147">[arXiv]</a></b> 
+        </p>
+        </div>
+        """)
+
+    with gr.Tabs():
+        with gr.TabItem("Canny"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Canny Edge Maps")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Canny', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        low_threshold = gr.Slider(label="Canny low threshold", minimum=1, maximum=255, value=40, step=1)
+                        high_threshold = gr.Slider(label="Canny high threshold", minimum=1, maximum=255, value=200,
+                                                   step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, low_threshold, high_threshold, condition_mode]
+            run_button.click(fn=process_canny, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("HED"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with HED Maps")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> HED', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="HED Resolution", minimum=128, maximum=1024, value=512,
+                                                      step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_hed, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Sketch"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Sketch Maps")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Sketch', value=False)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="HED Resolution", minimum=128, maximum=1024, value=512,
+                                                      step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_sketch, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Depth"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Depth Maps")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Depth', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="Depth Resolution", minimum=128, maximum=1024, value=384,
+                                                      step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_depth, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Normal"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Normal Surface")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Normal', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="Depth Resolution", minimum=128, maximum=1024, value=384,
+                                                      step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_normal, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Human Pose"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Human Pose")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Skeleton', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="OpenPose Resolution", minimum=128, maximum=1024, value=512,
+                                                      step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_pose, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Segmentation"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Segmentation Maps (ADE20K)")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Seg', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        detect_resolution = gr.Slider(label="Segmentation Resolution", minimum=128, maximum=1024,
+                                                      value=512, step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution,
+                   ddim_steps, guess_mode, strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_seg, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Bbox"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Object Bounding Boxes (MS-COCO)")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Bbox', value=True)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        confidence = gr.Slider(label="Confidence of Detection", minimum=0.1, maximum=1.0, value=0.4,
+                                               step=0.1)
+                        nms_thresh = gr.Slider(label="Nms Threshold", minimum=0.1, maximum=1.0, value=0.5, step=0.1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, bright')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, confidence, nms_thresh, condition_mode]
+            run_button.click(fn=process_bbox, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Outpainting"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Image Outpainting")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: Extending', value=False)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+
+                        height_top_extended = gr.Slider(label="Top Extended Ratio (%)", minimum=1, maximum=200,
+                                                        value=50, step=1)
+                        height_down_extended = gr.Slider(label="Down Extended Ratio (%)", minimum=1, maximum=200,
+                                                         value=50, step=1)
+
+                        width_left_extended = gr.Slider(label="Left Extended Ratio (%)", minimum=1, maximum=200,
+                                                        value=50, step=1)
+                        width_right_extended = gr.Slider(label="Right Extended Ratio (%)", minimum=1, maximum=200,
+                                                         value=50, step=1)
+
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, height_top_extended, height_down_extended, width_left_extended, width_right_extended, condition_mode]
+            run_button.click(fn=process_outpainting, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Inpainting"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Image Inpainting")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: Cropped Masking', value=False)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        h_ratio_t = gr.Slider(label="Top Masking Ratio (%)", minimum=0, maximum=100, value=30,
+                                              step=1)
+                        h_ratio_d = gr.Slider(label="Down Masking Ratio (%)", minimum=0, maximum=100, value=60,
+                                              step=1)
+                        w_ratio_l = gr.Slider(label="Left Masking Ratio (%)", minimum=0, maximum=100, value=30,
+                                              step=1)
+                        w_ratio_r = gr.Slider(label="Right Masking Ratio (%)", minimum=0, maximum=100, value=60,
+                                              step=1)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, h_ratio_t, h_ratio_d, w_ratio_l, w_ratio_r, condition_mode]
+            run_button.click(fn=process_inpainting, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Colorization"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Gray Image Colorization")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Gray', value=False)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, colorful')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, condition_mode]
+            run_button.click(fn=process_colorization, inputs=ips, outputs=[result_gallery])
+
+        with gr.TabItem("Deblurring"):
+            with gr.Row():
+                gr.Markdown("## UniControl Stable Diffusion with Image Deblurring")
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', type="numpy")
+                    prompt = gr.Textbox(label="Prompt")
+                    run_button = gr.Button(label="Run")
+                    with gr.Accordion("Advanced options", open=False):
+                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                        image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512,
+                                                     step=64)
+                        strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                        condition_mode = gr.Checkbox(label='Condition Extraction: RGB -> Blur', value=False)
+                        guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                        ksize = gr.Slider(label="Kernel Size", minimum=11, maximum=101, value=51, step=2)
+                        ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                        eta = gr.Number(label="eta (DDIM)", value=0.0)
+                        a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
+                        n_prompt = gr.Textbox(label="Negative Prompt", value='')
+                with gr.Column():
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2,
+                                                                                                           height='auto')
+            ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode,
+                   strength, scale, seed, eta, ksize, condition_mode]
+            run_button.click(fn=process_deblur, inputs=ips, outputs=[result_gallery])
+
+
+    gr.Markdown('''### Tips
+         -  Please pay attention to <u> Condition Extraction </u> option. 
+         - Positive prompts and negative prompts are very useful sometimes. 
+        ''')
+    gr.Markdown('''### Related Spaces
+    - https://huggingface.co/spaces/hysts/ControlNet
+    - https://huggingface.co/spaces/shi-labs/Prompt-Free-Diffusion
+    ''')
+demo.launch()

annotator/blur/__init__.py

@@ -0,0 +1,7 @@
+import cv2
+
+class Blurrer:
+    def __call__(self, img, ksize):
+        img_new = cv2.GaussianBlur(img, (ksize, ksize), cv2.BORDER_DEFAULT)
+        img_new = img_new.astype('ubyte')
+        return img_new

annotator/canny/__init__.py

@@ -0,0 +1,16 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+import cv2
+
+
+class CannyDetector:
+    def __call__(self, img, low_threshold, high_threshold):
+        return cv2.Canny(img, low_threshold, high_threshold)

annotator/ckpts/ckpts.txt

@@ -0,0 +1 @@
+Weights here.

annotator/grayscale/__init__.py

@@ -0,0 +1,5 @@
+from skimage import color
+
+class GrayscaleConverter:
+    def __call__(self, img):
+        return (color.rgb2gray(img) * 255.0).astype('ubyte')

annotator/hed/__init__.py

@@ -0,0 +1,107 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+# This is an improved version and model of HED edge detection without GPL contamination
+# Please use this implementation in your products
+# This implementation may produce slightly different results from Saining Xie's official implementations,
+# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
+# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
+# and in this way it works better for gradio's RGB protocol
+
+import os
+import cv2
+import torch
+import numpy as np
+
+from einops import rearrange
+from annotator.util import annotator_ckpts_path
+
+
+class DoubleConvBlock(torch.nn.Module):
+    def __init__(self, input_channel, output_channel, layer_number):
+        super().__init__()
+        self.convs = torch.nn.Sequential()
+        self.convs.append(torch.nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        for i in range(1, layer_number):
+            self.convs.append(torch.nn.Conv2d(in_channels=output_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        self.projection = torch.nn.Conv2d(in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0)
+
+    def __call__(self, x, down_sampling=False):
+        h = x
+        if down_sampling:
+            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
+        for conv in self.convs:
+            h = conv(h)
+            h = torch.nn.functional.relu(h)
+        return h, self.projection(h)
+
+
+class ControlNetHED_Apache2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
+        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
+        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
+        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
+        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
+        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
+
+    def __call__(self, x):
+        h = x - self.norm
+        h, projection1 = self.block1(h)
+        h, projection2 = self.block2(h, down_sampling=True)
+        h, projection3 = self.block3(h, down_sampling=True)
+        h, projection4 = self.block4(h, down_sampling=True)
+        h, projection5 = self.block5(h, down_sampling=True)
+        return projection1, projection2, projection3, projection4, projection5
+
+
+class HEDdetector:
+    def __init__(self):
+        remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth"
+        modelpath = remote_model_path
+        modelpath = os.path.join(annotator_ckpts_path, "ControlNetHED.pth")
+        if not os.path.exists(modelpath):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+        self.netNetwork = ControlNetHED_Apache2().float().cuda().eval()
+        self.netNetwork.load_state_dict(torch.load(modelpath))
+
+    def __call__(self, input_image):
+        assert input_image.ndim == 3
+        H, W, C = input_image.shape
+        with torch.no_grad():
+            image_hed = torch.from_numpy(input_image.copy()).float().cuda()
+            image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
+            edges = self.netNetwork(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+            return edge
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z

annotator/inpainting/__init__.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+class Inpainter:
+    def __call__(self, img, height_top_mask, height_down_mask, width_left_mask, width_right_mask):
+        h = img.shape[0]
+        w = img.shape[1]
+        h_top_mask = int(float(h) / 100.0 * float(height_top_mask))
+        h_down_mask = int(float(h) / 100.0 * float(height_down_mask))
+
+        w_left_mask = int(float(w) / 100.0 * float(width_left_mask))
+        w_right_mask = int(float(w) / 100.0 * float(width_right_mask))
+
+        img_new = img
+        img_new[h_top_mask:h_down_mask, w_left_mask:w_right_mask] = 0
+        img_new = img_new.astype('ubyte')
+        return img_new

annotator/midas/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Intel ISL (Intel Intelligent Systems Lab)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

annotator/midas/__init__.py

@@ -0,0 +1,52 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+# Midas Depth Estimation
+# From https://github.com/isl-org/MiDaS
+# MIT LICENSE
+
+import cv2
+import numpy as np
+import torch
+
+from einops import rearrange
+from .api import MiDaSInference
+
+
+class MidasDetector:
+    def __init__(self):
+        self.model = MiDaSInference(model_type="dpt_hybrid").cuda()
+
+    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):
+        assert input_image.ndim == 3
+        image_depth = input_image
+        with torch.no_grad():
+            image_depth = torch.from_numpy(image_depth).float().cuda()
+            image_depth = image_depth / 127.5 - 1.0
+            image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
+            depth = self.model(image_depth)[0]
+
+            depth_pt = depth.clone()
+            depth_pt -= torch.min(depth_pt)
+            depth_pt /= torch.max(depth_pt)
+            depth_pt = depth_pt.cpu().numpy()
+            depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
+
+            depth_np = depth.cpu().numpy()
+            x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3)
+            y = cv2.Sobel(depth_np, cv2.CV_32F, 0, 1, ksize=3)
+            z = np.ones_like(x) * a
+            x[depth_pt < bg_th] = 0
+            y[depth_pt < bg_th] = 0
+            normal = np.stack([x, y, z], axis=2)
+            normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
+            normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
+
+            return depth_image, normal_image

annotator/midas/api.py

@@ -0,0 +1,183 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+# based on https://github.com/isl-org/MiDaS
+
+import cv2
+import os
+import torch
+import torch.nn as nn
+from torchvision.transforms import Compose
+
+from .midas.dpt_depth import DPTDepthModel
+from .midas.midas_net import MidasNet
+from .midas.midas_net_custom import MidasNet_small
+from .midas.transforms import Resize, NormalizeImage, PrepareForNet
+from annotator.util import annotator_ckpts_path
+
+
+ISL_PATHS = {
+    "dpt_large": os.path.join(annotator_ckpts_path, "dpt_large_384.pt"),
+    "dpt_hybrid": os.path.join(annotator_ckpts_path, "dpt_hybrid-midas-501f0c75.pt"),
+    "midas_v21": "",
+    "midas_v21_small": "",
+}
+
+remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
+# remote_model_path = "https://storage.googleapis.com/sfr-unicontrol-data-research/annotator/ckpts/dpt_large_384.pt"  #"https://huggingface.co/Salesforce/UniControl/blob/main/annotator/ckpts/dpt_large_384.pt"
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def load_midas_transform(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load transform only
+    if model_type == "dpt_large":  # DPT-Large
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    elif model_type == "midas_v21_small":
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    else:
+        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return transform
+
+
+def load_model(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load network
+    model_path = ISL_PATHS[model_type]
+    if model_type == "dpt_large":  # DPT-Large
+        if not os.path.exists(model_path):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+        #model_path = remote_model_path
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitl16_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        if not os.path.exists(model_path):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
+
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitb_rn50_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        model = MidasNet(model_path, non_negative=True)
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    elif model_type == "midas_v21_small":
+        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
+                               non_negative=True, blocks={'expand': True})
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    else:
+        print(f"model_type '{model_type}' not implemented, use: --model_type large")
+        assert False
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return model.eval(), transform
+
+
+class MiDaSInference(nn.Module):
+    MODEL_TYPES_TORCH_HUB = [
+        "DPT_Large",
+        "DPT_Hybrid",
+        "MiDaS_small"
+    ]
+    MODEL_TYPES_ISL = [
+        "dpt_large",
+        "dpt_hybrid",
+        "midas_v21",
+        "midas_v21_small",
+    ]
+
+    def __init__(self, model_type):
+        super().__init__()
+        assert (model_type in self.MODEL_TYPES_ISL)
+        model, _ = load_model(model_type)
+        self.model = model
+        self.model.train = disabled_train
+
+    def forward(self, x):
+        with torch.no_grad():
+            prediction = self.model(x)
+        return prediction
+

annotator/midas/midas/base_model.py

@@ -0,0 +1,26 @@
+'''
+ * Copyright (c) 2023 Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: Apache License 2.0
+ * For full license text, see LICENSE.txt file in the repo root or http://www.apache.org/licenses/
+ * By Can Qin
+ * Modified from ControlNet repo: https://github.com/lllyasviel/ControlNet
+ * Copyright (c) 2023 Lvmin Zhang and Maneesh Agrawala
+'''
+
+import torch
+
+
+class BaseModel(torch.nn.Module):
+    def load(self, path):
+        """Load model from file.
+
+        Args:
+            path (str): file path
+        """
+        parameters = torch.load(path, map_location=torch.device('cpu'))
+
+        if "optimizer" in parameters:
+            parameters = parameters["model"]
+
+        self.load_state_dict(parameters)