Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for SD 2.1 Turbo #14170

Merged
merged 2 commits into from
Dec 2, 2023
Merged

Conversation

MrCheeze
Copy link
Contributor

@MrCheeze MrCheeze commented Dec 2, 2023

Description

This adds support for the newly released Stable Diffusion 2.1 Turbo model, which is more lightweight than sdxl-turbo.

Until now, all SD1 and SD2 models have been released in the LDM format, and all SDXL models have been released in SGM format. Because of that, there is a strong assumption all over the webui codebase that SGM format implies SDXL.

However, the SD 2.1 Turbo checkpoint was released in SGM format, breaking that assumption. Rather than rewrite the whole codebase to account for this possibility (possibly breaking extensions), I opted just to convert the checkpoint from SGM to LDM on load. (Which is as simple as renaming the dict keys from conditioner.embedders.0.model.* to cond_stage_model.model.*). By doing this, all features and extensions that were working with SD2, immediately work with SD2 Turbo as well.

This PR also fixes an existing issue where undo_hijack was only partially undoing the hijack, causing loading of SD2 models to fail (because is_using_v_parameterization_for_sd2 was not able to run).

Screenshots/videos:

image

Checklist:

@AUTOMATIC1111 AUTOMATIC1111 merged commit 4125552 into AUTOMATIC1111:dev Dec 2, 2023
3 checks passed
w-e-w added a commit that referenced this pull request Dec 2, 2023
@w-e-w w-e-w mentioned this pull request Dec 16, 2023
@light-and-ray
Copy link
Contributor

Nice, works well with StableSR! But I have the bug with model switching with this model too #13763 I have to provide this model via --ckpt argument

Maybe you have any idea, why it happens?

Loading model sd_v2-1_768-ema-pruned.safetensors [30970b314c] (3 out of 5)
Loading weights [30970b314c] from /home/user/ai-apps/stable-diffusion-webui/models/Stable-diffusion/sd_v2-1_768-ema-pruned.safetensors
changing setting sd_model_checkpoint to sd_v2-1_768-ema-pruned.safetensors [30970b314c]: NotImplementedError
Traceback (most recent call last):
  File "/home/user/ai-apps/stable-diffusion-webui/modules/options.py", line 146, in set
    option.onchange()
  File "/home/user/ai-apps/stable-diffusion-webui/modules/call_queue.py", line 13, in f
    res = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/initialize_util.py", line 174, in <lambda>
    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: sd_models.reload_model_weights()), call=False)
                                                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models.py", line 764, in reload_model_weights
    sd_model = reuse_model_from_already_loaded(sd_model, checkpoint_info, timer)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models.py", line 733, in reuse_model_from_already_loaded
    load_model(checkpoint_info)
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models.py", line 618, in load_model
    checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models_config.py", line 119, in find_checkpoint_config
    return guess_model_config_from_state_dict(state_dict, info.filename)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models_config.py", line 91, in guess_model_config_from_state_dict
    elif is_using_v_parameterization_for_sd2(sd):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_models_config.py", line 64, in is_using_v_parameterization_for_sd2
    out = (unet(x_test, torch.asarray([999], device=device), context=test_cond) - x_test).mean().item()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_unet.py", line 91, in UNetModel_forward
    return original_forward(self, x, timesteps, context, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/openaimodel.py", line 797, in forward
    h = module(h, emb, context)
        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/openaimodel.py", line 84, in forward
    x = layer(x, context)
        ^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/attention.py", line 334, in forward
    x = block(x, context=context[i])
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/attention.py", line 269, in forward
    return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/util.py", line 121, in checkpoint
    return CheckpointFunction.apply(func, len(inputs), *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/util.py", line 136, in forward
    output_tensors = ctx.run_function(*ctx.input_tensors)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/modules/attention.py", line 272, in _forward
    x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/modules/sd_hijack_optimizations.py", line 496, in xformers_attention_forward
    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/xformers/ops/fmha/__init__.py", line 192, in memory_efficient_attention
    return _memory_efficient_attention(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/xformers/ops/fmha/__init__.py", line 290, in _memory_efficient_attention
    return _memory_efficient_attention_forward(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/xformers/ops/fmha/__init__.py", line 306, in _memory_efficient_attention_forward
    op = _dispatch_fw(inp)
         ^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/xformers/ops/fmha/dispatch.py", line 94, in _dispatch_fw
    return _run_priority_list(
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user/ai-apps/stable-diffusion-webui/venv/lib/python3.11/site-packages/xformers/ops/fmha/dispatch.py", line 69, in _run_priority_list
    raise NotImplementedError(msg)
NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:

The problem only with sd2.1 and only on reload

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants