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

Support LinFusion #9596

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Support LinFusion #9596

wants to merge 9 commits into from

Conversation

Huage001
Copy link

@Huage001 Huage001 commented Oct 6, 2024

What does this PR do?

Support LinFusion. It accelerates diffusion models by replacing all the self-attention layers in a diffusion UNet with distilled Generalized Linear Attention layers. The distilled model is linear-complexity and highly compatible with existing diffusion plugins like ControlNet, IP-Adapter, LoRA, etc. The acceleration can be dramatic at high resolution. Strategical pipelines for high-resolution generation can be found in the original codebase.

You can use it with only 1 additional line:

import torch
from diffusers import StableDiffusionPipeline

repo_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16, variant="fp16").to("cuda")

+ pipe.load_linfusion(pipeline_name_or_path=repo_id)

image = pipe("a photo of an astronaut on a moon").images[0]

Currently, stable-diffusion-v1-5/stable-diffusion-v1-5, stabilityai/stable-diffusion-2-1, stabilityai/stable-diffusion-xl-base-1.0, models finetuned from them, and pipelines based on them are supported. If the repo_id is different from them, e.g., when using a fine-tuned model from the community, you need to specify pipeline_name_or_path explicitly to the model it is based on. Otherwise, this argument is optional and LinFusion will read it from the current pipeline. Alternatively, you can also specify the argument pretrained_model_name_or_path_or_dict to load LinFusion from other sources. You can also unload it with pipe.unload_linfusion() when unnecessary.

Accordingly, we also update the doc under docs/source/en/optimization/linfusion.md for a specific example.

Thanks for your efforts in reviewing this pull request in advance! We are open to any changes to make sure LinFusion can best fit the current diffusers library!

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

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.

2 participants