[Issue] ControlNet Support

[paulo-coronado]

ToMe is currently not working with ControlNet. I tried canny preprocessor, for example, and got the following error:

TypeError: cannot unpack non-iterable NoneType object

Now, I am not sure if sd-web-ui-controlnet doesn't support ToMe or the opposite... 🤔

[AmanKishore]

Getting the same error!

[dbolya]

ControlNet might patch the same things that ToMe patches, causing a conflict.
Can someone find which line is causing that TypeError? I haven't used ControlNet before.

[AmanKishore]

Stack trace:

  File "/root/sag/pipeline_stable_diffusion_controlnet.py", line 947, in __call__
    noise_pred = self.unet(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/diffusers/models/unet_2d_condition.py", line 620, in forward
    sample, res_samples = downsample_block(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/diffusers/models/unet_2d_blocks.py", line 837, in forward
    hidden_states = attn(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/diffusers/models/transformer_2d.py", line 265, in forward
    hidden_states = block(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/tomesd/patch.py", line 84, in forward
    m_a, m_c, m_m, u_a, u_c, u_m = compute_merge(hidden_states, self._tome_info)
  File "/usr/local/lib/python3.9/site-packages/tomesd/patch.py", line 11, in compute_merge
    original_h, original_w = tome_info["size"]
TypeError: cannot unpack non-iterable NoneType object

[dbolya]

Hmm, seems like the controlnet plugin might overwrite the forward function for the unet itself, which is how ToMe gets the image size. I'll try to see if I can find a workaround.

[dbolya]

Fixed this by changing how I patched the model. However, do note that controlnet runs its own model in parallel, and that model doesn't have tome applied. So the speed-up might not be as good as without using controlnet.

If there are any more issues, feel free to reopen.

[icech]

@dbolya
do you plan to add the feature of ToMe for speeding up the ControlNet

[papuSpartan]

@icech
It's up to the extension author whether they will add it, It's not a problem with tomesd.