Repository: stable-diffusion-webui
Issue: [Bug]: txt2img fails when using xformers on Google Colab
Labels: asking-for-help-with-local-system-issues

### Is there an existing issue for this?

- [X] I have searched the existing issues and checked the recent builds/commits

### What happened?

txt2img fails when using xformers on Google Colab. 

- Using xformers from pip (0.0.12) will give TypeError: memory_efficient_attention() got an unexpected keyword argument 'attn_bias'
- Using xformers wheel from naifu (0.0.14) gives RuntimeError: Expected query.dim() == 3 to be true, but got false.
- Not installing triton or installing it without --pre produces some warning, but the error on txt2img stays the same.
- Webui won't automatically install xformers on Google Colab

### Steps to reproduce the problem

1. Install triton with --pre. The one from pip is old (1.1.1) and produces warning.
2. Install xformers
a. Installing from pip will install xformers 0.012
b. There's 0.014 from naifu colab: 
https://github.com/metrolobo/xformers_wheels/releases/download/1d31a3ac_various_6/xformers-0.0.14.dev0-cp37-cp37m-linux_x86_64.whl
3. Load NAI model (with or without vae doesn't matter).
4. Use txt2img. Here I'm using "ganyu \(genshin impact \)"

### What should have happened?

It generates an image.

### Commit where the problem happens

7f8ab1e

### What platforms do you use to access UI ?

Other/Cloud

### What browsers do you use to access the UI ?

Google Chrome

### Command Line Arguments

```Shell
--autolaunch --ckpt {workspace3}/model.ckpt --xformers
--autolaunch --ckpt {workspace3}/model.ckpt --vae-path {workspace3}/model.vae.pt --xformers
```


### Additional information, context and logs


<details>
<summary>Using xformers from pip will give TypeError: memory_efficient_attention() got an unexpected keyword argument 'attn_bias'</summary>

```
Running on local URL:  http://127.0.0.1:7860/

To create a public link, set share=True in launch().
  0% 0/20 [00:03<?, ?it/s]
Error completing request
Arguments: ('ganyu \\(genshin impact\\), ', 'guan yu', 'None', 'None', 20, 0, False, False, 1, 1, 7, -1.0, -1.0, 0, 0, 0, False, 512, 512, False, 0.7, 0, 0, 0, False, None, '', 'outputs', '', 1, True, 100, False, '<p style="margin-bottom:0.75em">Keyframe Format: <br>Seed | Prompt or just Prompt</p>', '', 25, True, 5.0, False, None, '', 10.0, True, 30.0, True, False, 1, '', 0, '', True, False, False) {}
Traceback (most recent call last):
  File "/content/nai/stable-diffusion-webui/modules/ui.py", line 212, in f
    res = list(func(*args, **kwargs))
  File "webui.py", line 63, in f
    res = func(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/txt2img.py", line 44, in txt2img
    processed = process_images(p)
  File "/content/nai/stable-diffusion-webui/modules/processing.py", line 411, in process_images
    samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength)
  File "/content/nai/stable-diffusion-webui/modules/processing.py", line 549, in sample
    samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning)
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 417, in sample
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 326, in launch_sampling
    return func()
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 417, in <lambda>
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/sampling.py", line 80, in sample_euler_ancestral
    denoised = model(x, sigmas[i] * s_in, **extra_args)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 248, in forward
    x_out = self.inner_model(x_in, sigma_in, cond=cond_in)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/external.py", line 112, in forward
    eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/external.py", line 138, in get_eps
    return self.inner_model.apply_model(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/models/diffusion/ddpm.py", line 987, in apply_model
    x_recon = self.model(x_noisy, t, **cond)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/models/diffusion/ddpm.py", line 1410, in forward
    out = self.diffusion_model(x, t, context=cc)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/openaimodel.py", line 732, in forward
    h = module(h, emb, context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/openaimodel.py", line 85, in forward
    x = layer(x, context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 258, in forward
    x = block(x, context=context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 209, in forward
    return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/util.py", line 114, in checkpoint
    return CheckpointFunction.apply(func, len(inputs), *args)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/util.py", line 127, in forward
    output_tensors = ctx.run_function(*ctx.input_tensors)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 212, in _forward
    x = self.attn1(self.norm1(x)) + x
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/sd_hijack_optimizations.py", line 227, in xformers_attention_forward
    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
TypeError: memory_efficient_attention() got an unexpected keyword argument 'attn_bias'
```

</details>

<details>
<summary>Using xformers wheel from naifu gives RuntimeError: Expected query.dim() == 3 to be true, but got false. </summary>

```
Traceback (most recent call last):
  File "/content/nai/stable-diffusion-webui/modules/ui.py", line 212, in f
    res = list(func(*args, **kwargs))
  File "webui.py", line 63, in f
    res = func(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/txt2img.py", line 44, in txt2img
    processed = process_images(p)
  File "/content/nai/stable-diffusion-webui/modules/processing.py", line 411, in process_images
    samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength)
  File "/content/nai/stable-diffusion-webui/modules/processing.py", line 549, in sample
    samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning)
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 417, in sample
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 326, in launch_sampling
    return func()
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 417, in <lambda>
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/sampling.py", line 80, in sample_euler_ancestral
    denoised = model(x, sigmas[i] * s_in, **extra_args)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/sd_samplers.py", line 248, in forward
    x_out = self.inner_model(x_in, sigma_in, cond=cond_in)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/external.py", line 112, in forward
    eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/k-diffusion/k_diffusion/external.py", line 138, in get_eps
    return self.inner_model.apply_model(*args, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/models/diffusion/ddpm.py", line 987, in apply_model
    x_recon = self.model(x_noisy, t, **cond)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/models/diffusion/ddpm.py", line 1410, in forward
    out = self.diffusion_model(x, t, context=cc)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/openaimodel.py", line 732, in forward
    h = module(h, emb, context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/openaimodel.py", line 85, in forward
    x = layer(x, context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 258, in forward
    x = block(x, context=context)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 209, in forward
    return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/util.py", line 114, in checkpoint
    return CheckpointFunction.apply(func, len(inputs), *args)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/diffusionmodules/util.py", line 127, in forward
    output_tensors = ctx.run_function(*ctx.input_tensors)
  File "/content/nai/stable-diffusion-webui/repositories/stable-diffusion/ldm/modules/attention.py", line 212, in _forward
    x = self.attn1(self.norm1(x)) + x
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/nai/stable-diffusion-webui/modules/sd_hijack_optimizations.py", line 227, in xformers_attention_forward
    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
  File "/usr/local/lib/python3.7/dist-packages/xformers/ops.py", line 575, in memory_efficient_attention
    query=query, key=key, value=value, attn_bias=attn_bias, p=p
  File "/usr/local/lib/python3.7/dist-packages/xformers/ops.py", line 196, in forward_no_grad
    causal=isinstance(attn_bias, LowerTriangularMask),
  File "/usr/local/lib/python3.7/dist-packages/torch/_ops.py", line 143, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: Expected query.dim() == 3 to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)
```

</details>

<details>
<summary>Not installing triton produces this message, but the error on txt2img is the same</summary>

```
WARNING:root:Triton is not available, some optimizations will not be enabled.
Error No module named 'triton'
```

</details>

<details>
<summary>Installing triton without --pre produces this message, but the error on txt2img is the same</summary>

```
WARNING:root:A matching Triton is not available, some optimizations will not be enabled.
Error caught was: module 'triton.language' has no attribute 'constexpr'
```

</details>

- Webui won't install xformers automatically on Google Colab.

Related: #2731