Repository navigation
LoRA weights baking. #4642
Description
Is your feature request related to a problem? Please describe.
LoRA is consuming extra GPU power that is not required in inference. This can cause a noticable amount of overhead especially with high resolutions. For higher LoRA ranks (64/128), this can be up to 30%-50% performance gain.
Describe the solution you'd like
A bake_lora_weights function on UNet models (or any model that supports attention processor).
Uses can decide whether they want to keep swapping LoRA weights available or to bake weights to reduce inference overhead.
Proof of Concept:
@torch.no_grad()
def bake_lora_weights(self):
def expansion(lora: LoRALinearLayer):
# in x in -> in x out
eye = torch.eye(lora.down.weight.size(1))
eye = eye.to(lora.down.weight)
return lora(eye).T
def traverse(name: str, module: Attention):
if hasattr(module, "set_processor"):
assert isinstance(module.processor, (LoRAAttnProcessor, LoRAAttnProcessor2_0, LoRAXFormersAttnProcessor))
module.to_k.weight += expansion(module.processor.to_k_lora)
module.to_v.weight += expansion(module.processor.to_v_lora)
module.to_q.weight += expansion(module.processor.to_q_lora)
module.to_out[0].weight += expansion(module.processor.to_out_lora)
for sub_name, child in module.named_children():
traverse(f"{name}.{sub_name}", child)
for name, module in self.named_children():
traverse(name, module)
if torch.__version__ >= '2.0':
self.unet.set_attn_processor(AttnProcessor2_0())
elif is_xformers_available():
self.unet.set_attn_processor(XFormersAttnProcessor())
else:
self.unet.set_attn_processor(AttnProcessor())Describe alternatives you've considered
N/A
Additional context
N/A
Activity
apolinario commentedon Aug 17, 2023
Check this PR out :-)
#4473
eliphatfs commentedon Aug 17, 2023
Sorry for duplication. I only searched for issues and did not search for PRs. Thanks for sharing that pointer.