add depth-anything controlnet

modules/control/units/controlnet.py

@@ -25,6 +25,7 @@ predefined_sd15 = {
     'Shuffle': "lllyasviel/control_v11e_sd15_shuffle",
     'SoftEdge': "lllyasviel/control_v11p_sd15_softedge",
     'Tile': "lllyasviel/control_v11f1e_sd15_tile",
+    'Depth Anything': 'vladmandic/depth-anything',
     'Canny FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_canny.safetensors',
     'Inpaint FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_inpaint.safetensors',
     'LineArt Anime FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_animeline.safetensors',
@@ -116,21 +117,29 @@ class ControlNet():
 
     def load_safetensors(self, model_path):
         name = os.path.splitext(model_path)[0]
-        yaml_path = None
+        config_path = None
         if not os.path.exists(model_path):
             import huggingface_hub as hf
             parts = model_path.split('/')
             repo_id = f'{parts[0]}/{parts[1]}'
             filename = os.path.splitext('/'.join(parts[2:]))[0]
             model_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.safetensors', cache_dir=cache_dir)
-            try:
-                yaml_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.yaml', cache_dir=cache_dir)
-            except Exception:
-                pass # no yaml file
+            if config_path is None:
+                try:
+                    config_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.yaml', cache_dir=cache_dir)
+                except Exception:
+                    pass # no yaml file
+            if config_path is None:
+                try:
+                    config_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.json', cache_dir=cache_dir)
+                except Exception:
+                    pass # no yaml file
         elif os.path.exists(name + '.yaml'):
-            yaml_path = f'{name}.yaml'
-        if yaml_path is not None:
-            self.load_config['original_config_file '] = yaml_path
+            config_path = f'{name}.yaml'
+        elif os.path.exists(name + '.json'):
+            config_path = f'{name}.json'
+        if config_path is not None:
+            self.load_config['original_config_file '] = config_path
         self.model = ControlNetModel.from_single_file(model_path, **self.load_config)
 
     def load(self, model_id: str = None) -> str:

modules/control/units/xs_model.py

@@ -26,16 +26,15 @@ from diffusers.models.attention_processor import USE_PEFT_BACKEND, AttentionProc
 from diffusers.models.autoencoders import AutoencoderKL
 from diffusers.models.lora import LoRACompatibleConv
 from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.unet_2d_blocks import (
-    CrossAttnDownBlock2D,
-    CrossAttnUpBlock2D,
-    DownBlock2D,
-    Downsample2D,
-    ResnetBlock2D,
-    Transformer2DModel,
-    UpBlock2D,
-    Upsample2D,
-)
+try:
+    from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, Downsample2D, ResnetBlock2D, Transformer2DModel, UpBlock2D, Upsample2D
+except Exception:
+    pass
+try:
+    from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, Downsample2D, ResnetBlock2D, Transformer2DModel, UpBlock2D, Upsample2D
+except Exception:
+    pass
+
 from diffusers.models.unet_2d_condition import UNet2DConditionModel
 from diffusers.utils import BaseOutput, logging
 

modules/lama.py

@@ -87,6 +87,12 @@ class SimpleLama:
         self.model.to(self.device)
 
     def __call__(self, image: Image.Image | np.ndarray, mask: Image.Image | np.ndarray):
+        if image is None:
+            log.warning('LaMa: image is none')
+            return None
+        if mask is None:
+            mask = Image.new('L', image.size, 0)
+            return None
         image, mask = prepare_img_and_mask(image, mask, self.device)
         with devices.inference_context():
             inpainted = self.model(image, mask)

modules/masking.py

@@ -217,6 +217,8 @@ def run_segment(input_image: gr.Image, input_mask: np.ndarray):
             continue
         overlap = 0
         if input_mask_size > 0:
+            if mask.shape != input_mask.shape:
+                mask = cv2.resize(mask, (input_mask.shape[1], input_mask.shape[0]), interpolation=cv2.INTER_CUBIC)
             overlap = cv2.bitwise_and(mask, input_mask)
             overlap = np.count_nonzero(overlap)
             if overlap == 0:

modules/ui_control.py

@@ -629,7 +629,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             settings.append(gr.Slider(label="Denoising steps", minimum=1, maximum=99, step=1, value=10))
                             settings.append(gr.Slider(label="Ensemble size", minimum=1, maximum=99, step=1, value=10))
                         with gr.Accordion('Depth Anything', open=True, elem_classes=['processor-settings']):
-                            settings.append(gr.Dropdown(label="Color map", choices=['inferno'] + masking.COLORMAP, value='inferno'))
+                            settings.append(gr.Dropdown(label="Color map", choices=['none'] + masking.COLORMAP, value='inferno'))
                         for setting in settings:
                             setting.change(fn=processors.update_settings, inputs=settings, outputs=[])