Repository: AUTOMATIC1111/stable-diffusion-webui

Issue:
[Bug]: Cannot run "512-depth-ema.ckpt" model
### Is there an existing issue for this?

- [X] I have searched the existing issues and checked the recent builds/commits

### What happened?

I want to run the depth model (512-depth-ema.ckpt (d522022)) of stable diffusion 2.0, but it came to an error. 
Though I have run the basic model (768-v-ema.ckpt (2c02b20a)) successfully on previous, the depth one yet not. I emphasize that I put the "512-depth-ema.yaml", which is renamed from "v2-inferance.yaml", correctly in /models/stable-diffusion/, next to the model. 
Puzzled, I redownload it, nothing different happen. I stress once again that I have run the basic one smoothly before!

### Steps to reproduce the problem

1. Add "512-depth-ema.yaml" and "512-depth-ema.yaml" to /models/stable-diffusion/
2. Go to webui-user.bat and run it
3. select the 512-depth-ema.ckpt (d0522d12) in WEBUI page
4. oops, it crash!!

### What should have happened?

I think no error should occur and I can produce depth image by send image in.

### Commit where the problem happens

in loading 512-depth-ema.ckpt, the commit hash is ce049c471b4a1d22f5a8fe8f527788edcf934eda

### What platforms do you use to access UI ?

Windows

### What browsers do you use to access the UI ?

Microsoft Edge

### Command Line Arguments

```Shell
Loading config from: D:\AI6\stable-diffusion-webui\models\Stable-diffusion\512-depth-ema.yaml
LatentDiffusion: Running in v-prediction mode
DiffusionWrapper has 865.91 M params.
Loading weights [d0522d12] from D:\AI6\stable-diffusion-webui\models\Stable-diffusion\512-depth-ema.ckpt
Traceback (most recent call last):
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\gradio\routes.py", line 284, in run_predict
    output = await app.blocks.process_api(
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 982, in process_api
    result = await self.call_function(fn_index, inputs, iterator)
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 824, in call_function
    prediction = await anyio.to_thread.run_sync(
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\anyio\to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "D:\AI6\stable-diffusion-webui\modules\ui.py", line 1616, in <lambda>
    fn=lambda value, k=k: run_settings_single(value, key=k),
  File "D:\AI6\stable-diffusion-webui\modules\ui.py", line 1457, in run_settings_single
    if not opts.set(key, value):
  File "D:\AI6\stable-diffusion-webui\modules\shared.py", line 478, in set
    self.data_labels[key].onchange()
  File "D:\AI6\stable-diffusion-webui\modules\call_queue.py", line 15, in f
    res = func(*args, **kwargs)
  File "D:\AI6\stable-diffusion-webui\webui.py", line 62, in <lambda>
    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights()))
  File "D:\AI6\stable-diffusion-webui\modules\sd_models.py", line 292, in reload_model_weights
    load_model(checkpoint_info)
  File "D:\AI6\stable-diffusion-webui\modules\sd_models.py", line 261, in load_model
    load_model_weights(sd_model, checkpoint_info)
  File "D:\AI6\stable-diffusion-webui\modules\sd_models.py", line 192, in load_model_weights
    model.load_state_dict(sd, strict=False)
  File "D:\AI6\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1604, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for LatentDiffusion:
        size mismatch for model.diffusion_model.input_blocks.0.0.weight: copying a param with shape torch.Size([320, 5, 3, 3]) from checkpoint, the shape in current model is torch.Size([320, 4, 3, 3]).
```


### Additional information, context and logs

The model hash is d522d022, and SHA1: a81e98c0f932a55d8dd08058c1b9ce68525ac6a5
The error seems coming from bad torch size setting.


Discussion:
Version 2.0 works differently from 1.x versions, and right now only the 768-v-ema.ckpt model is supported 
> Version 2.0 works differently from 1.x versions, and right now only the 768-v-ema.ckpt model is supported

Thanks, guy! I have replaced it by extension "depthmap2mask" [Extraltodeus](https://github.com/Extraltodeus).
However, I still want to use it in future...:-(
> Version 2.0 works differently from 1.x versions, and right now only the 768-v-ema.ckpt model is supported

Now! Boy! I have found a leak of yaml, which can run the 512-basic-ema.ckpt smoothly with webui.
[v2-inference-e.zip](https://github.com/AUTOMATIC1111/stable-diffusion-webui/files/10148302/v2-inference-e.zip)

> I have found a leak of yaml, which can run the 512-basic-ema.ckpt smoothly with webui.

Thanks @LieDeath. This yaml config file worked for me. 

I was not able to load 512-basic-ema with original instructions.

Where is that file coming from if I may ask? 
@healthyfat Did you download the right one from the original instructions? They're currently confusingly-phrased, with the location of the 512 file tacked on as an afterthought.

![image](https://user-images.githubusercontent.com/118340091/206285490-548829e0-1da2-4ccc-a0d1-0f5f6416c9f3.png)

(Amusingly: All of these files get downloaded for us to `repositories\stable-diffusion-stability-ai\configs\stable-diffusion` anyway when the SD2.0 repo is cloned into a subfolder; it's just a matter of copying them to a place that this repo looks for them)
@AnonymousCervine 

> Did you download the right one from the original instructions?

Yes. I used the config from the link you highlighted. It doesn't work for me. 
Only v2-inference-e.yaml does. 
Interesting. I wonder what the difference between our setups is.

(Probably as much as should be spoken of here; it's off-topic to this issue and probably needs its own?)
> Interesting. I wonder what the difference between our setups is.
> 
> (Probably as much as should be spoken of here; it's off-topic to this issue and probably needs its own?)

I take a review and find the only difference between **v2-inference-e.yaml** and **v2-inference.yaml** is _a line break_ at the end of the config file. Perhaps on some devices the missing line break in the **v2-inference.yaml** could cause problems:-(.

Resolving patch:
diff --git a/modules/processing.py b/modules/processing.py
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -21,7 +21,10 @@
 import modules.images as images
 import modules.styles
 import logging
+from ldm.data.util import AddMiDaS
+from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion
 
+from einops import repeat, rearrange
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
@@ -150,11 +153,26 @@ def txt2img_image_conditioning(self, x, width=None, height=None):
 
         return image_conditioning
 
-    def img2img_image_conditioning(self, source_image, latent_image, image_mask = None):
-        if self.sampler.conditioning_key not in {'hybrid', 'concat'}:
-            # Dummy zero conditioning if we're not using inpainting model.
-            return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)
+    def depth2img_image_conditioning(self, source_image):
+        # Use the AddMiDaS helper to Format our source image to suit the MiDaS model
+        transformer = AddMiDaS(model_type="dpt_hybrid")
+        transformed = transformer({"jpg": rearrange(source_image[0], "c h w -> h w c")})
+        midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)
+        midas_in = repeat(midas_in, "1 ... -> n ...", n=self.batch_size)
+
+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))
+        conditioning = torch.nn.functional.interpolate(
+            self.sd_model.depth_model(midas_in),
+            size=conditioning_image.shape[2:],
+            mode="bicubic",
+            align_corners=False,
+        )
+
+        (depth_min, depth_max) = torch.aminmax(conditioning)
+        conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.
+        return conditioning
 
+    def inpainting_image_conditioning(self, source_image, latent_image, image_mask = None):
         self.is_using_inpainting_conditioning = True
 
         # Handle the different mask inputs
@@ -191,6 +209,18 @@ def img2img_image_conditioning(self, source_image, latent_image, image_mask = No
 
         return image_conditioning
 
+    def img2img_image_conditioning(self, source_image, latent_image, image_mask=None):
+        # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely
+        # identify itself with a field common to all models. The conditioning_key is also hybrid.
+        if isinstance(self.sd_model, LatentDepth2ImageDiffusion):
+            return self.depth2img_image_conditioning(source_image)
+
+        if self.sampler.conditioning_key in {'hybrid', 'concat'}:
+            return self.inpainting_image_conditioning(source_image, latent_image, image_mask=image_mask)
+
+        # Dummy zero conditioning if we're not using inpainting or depth model.
+        return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)
+
     def init(self, all_prompts, all_seeds, all_subseeds):
         pass
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -7,6 +7,9 @@
 import re
 import safetensors.torch
 from omegaconf import OmegaConf
+from os import mkdir
+from urllib import request
+import ldm.modules.midas as midas
 
 from ldm.util import instantiate_from_config
 
@@ -36,6 +39,7 @@ def setup_model():
         os.makedirs(model_path)
 
     list_models()
+    enable_midas_autodownload()
 
 
 def checkpoint_tiles(): 
@@ -227,6 +231,48 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     sd_vae.load_vae(model, vae_file)
 
 
+def enable_midas_autodownload():
+    """
+    Gives the ldm.modules.midas.api.load_model function automatic downloading.
+
+    When the 512-depth-ema model, and other future models like it, is loaded,
+    it calls midas.api.load_model to load the associated midas depth model.
+    This function applies a wrapper to download the model to the correct
+    location automatically.
+    """
+
+    midas_path = os.path.join(models_path, 'midas')
+
+    # stable-diffusion-stability-ai hard-codes the midas model path to
+    # a location that differs from where other scripts using this model look.
+    # HACK: Overriding the path here.
+    for k, v in midas.api.ISL_PATHS.items():
+        file_name = os.path.basename(v)
+        midas.api.ISL_PATHS[k] = os.path.join(midas_path, file_name)
+
+    midas_urls = {
+        "dpt_large": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",
+        "dpt_hybrid": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt",
+        "midas_v21": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt",
+        "midas_v21_small": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt",
+    }
+
+    midas.api.load_model_inner = midas.api.load_model
+
+    def load_model_wrapper(model_type):
+        path = midas.api.ISL_PATHS[model_type]
+        if not os.path.exists(path):
+            if not os.path.exists(midas_path):
+                mkdir(midas_path)
+    
+            print(f"Downloading midas model weights for {model_type} to {path}")
+            request.urlretrieve(midas_urls[model_type], path)
+            print(f"{model_type} downloaded")
+
+        return midas.api.load_model_inner(model_type)
+
+    midas.api.load_model = load_model_wrapper
+
 def load_model(checkpoint_info=None):
     from modules import lowvram, sd_hijack
     checkpoint_info = checkpoint_info or select_checkpoint()


Test patch: