How to remove lora after loading it?

[alexblattner]

The title pretty much. I am trying to load and unload inside the denoising loop to apply the loras to specific parts of the image only

[pcuenca]

Hi @alexblattner, you could try pipe.unet.set_default_attn_processor(), would that work for your use case?

[alexblattner]

I'll try that, but I need to fix a lot of other bugs first. thanks @pcuenca

[pcuenca]

Sure :) In that case I'll be closing this issue, feel free to reopen if needed when you try it out :)

[alexblattner]

@pcuenca I am not sure if it worked or not. I may have related or unrelated issues. here's the breakdown:
in the denoising part, I have this code:

latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
               latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
               # predict the noise residual
               noise_preds = []
               for i in range(len(mask_list)):
                   if loras_apply is not None and loras_apply.get(i) is not None:
                       # self.load_lora_weights(loras_apply.get(i))
                       load_lora_weights(self, loras_apply.get(i), 1,self.device.type,dtype)
                   noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings[i]).sample
                   noise_preds.append(noise_pred)
                   if loras_apply is not None and loras_apply.get(i) is not None:
                       self.unet.set_default_attn_processor()
               # perform guidance
               if do_classifier_free_guidance:
                   noise_pred_unconds = []
                   noise_pred_texts = []
                   for i in range(len(mask_list)):
                     noise_pred_uncond, noise_pred_text = noise_preds[i].chunk(2)
                     noise_pred_unconds.append(noise_pred_uncond)
                     noise_pred_texts.append(noise_pred_text)
                   
                   result = None
                   noise_preds = []
                   for i in range(len(mask_list)):
                     noise_pred = noise_pred_unconds[i] + guidance_scale * (noise_pred_texts[i] - noise_pred_unconds[i])
                     noise_preds.append(noise_pred)
                   result = noise_preds[0] * mask_list[0]
                   for i in range(1, len(mask_list)):
                     result += noise_preds[i] * mask_list[i]

                   #noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                   # compute the previous noisy sample x_t -> x_t-1
                   latents = self.scheduler.step(result, t, latents, **extra_step_kwargs).prev_sample

                   # call the callback, if provided
                   if callback is not None and i % callback_steps == 0:
                       callback(i, t, latents)

the important part is this one:

if loras_apply is not None and loras_apply.get(i) is not None:
                        load_lora_weights(self, loras_apply.get(i), 1,self.device.type,dtype
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings[i]).sample
noise_preds.append(noise_pred)
if loras_apply is not None and loras_apply.get(i) is not None:
                        self.unet.set_default_attn_processor()

I am using the load_lora_weights function and pass strings representing the location of safetensors files containing loras. it looks like this:

def load_lora_weights(pipeline, checkpoint_path, multiplier, device, dtype):
    LORA_PREFIX_UNET = "lora_unet"
    LORA_PREFIX_TEXT_ENCODER = "lora_te"
    # load LoRA weight from .safetensors
    if isinstance(checkpoint_path, str):

        state_dict = load_file(checkpoint_path, device=device)

        updates = defaultdict(dict)
        for key, value in state_dict.items():
            # it is suggested to print out the key, it usually will be something like below
            # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

            layer, elem = key.split('.', 1)
            updates[layer][elem] = value

        # directly update weight in diffusers model
        for layer, elems in updates.items():

            if "text" in layer:
                layer_infos = layer.split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
                curr_layer = pipeline.text_encoder
            else:
                layer_infos = layer.split(LORA_PREFIX_UNET + "_")[-1].split("_")
                curr_layer = pipeline.unet

            # find the target layer
            temp_name = layer_infos.pop(0)
            while len(layer_infos) > -1:
                try:
                    curr_layer = curr_layer.__getattr__(temp_name)
                    if len(layer_infos) > 0:
                        temp_name = layer_infos.pop(0)
                    elif len(layer_infos) == 0:
                        break
                except Exception:
                    if len(temp_name) > 0:
                        temp_name += "_" + layer_infos.pop(0)
                    else:
                        temp_name = layer_infos.pop(0)

            # get elements for this layer
            weight_up = elems['lora_up.weight'].to(dtype)
            weight_down = elems['lora_down.weight'].to(dtype)
            alpha = elems['alpha']
            if alpha:
                alpha = alpha.item() / weight_up.shape[1]
            else:
                alpha = 1.0

            # update weight
            if len(weight_up.shape) == 4:
                curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up.squeeze(3).squeeze(2), weight_down.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
            else:
                curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up, weight_down)
    else:
        for ckptpath in checkpoint_path:
            pipeline=load_lora_weights(pipeline, ckptpath, multiplier, device, dtype)
    return pipeline

the error I am having is the following:
/mecomics-api/tester.py:33 in │
│ │
│ 30 buffer.seek(0) │
│ 31 image_bytes = buffer.read() │
│ 32 images = Image.open(BytesIO(image_bytes)) │
│ ❱ 33 image = pipe( │
│ 34 │ prompt=None, │
│ 35 │ negative_prompt=None, │
│ 36 │ prompt_embeds=promptE, │
│ │
│ /usr/local/lib/python3.8/site-packages/torch/utils/_contextlib.py:115 in decorate_context │
│ │
│ 112 │ @functools.wraps(func) │
│ 113 │ def decorate_context(*args, **kwargs): │
│ 114 │ │ with ctx_factory(): │
│ ❱ 115 │ │ │ return func(*args, **kwargs) │
│ 116 │ │
│ 117 │ return decorate_context │
│ 118 │
│ │
│ /mecomics-api/multilatents.py:986 in call │
│ │
│ 983 │ │ │ │ │ print(i) │
│ 984 │ │ │ │ │ if loras_apply is not None and loras_apply.get(i) is not None: │
│ 985 │ │ │ │ │ │ # self.load_lora_weights(loras_apply.get(i)) │
│ ❱ 986 │ │ │ │ │ │ load_lora_weights(self, loras_apply.get(i), 1,self.device.type,d │
│ 987 │ │ │ │ │ noise_pred=None │
│ 988 │ │ │ │ │ if controlnet_image is not None: │
│ 989 │ │ │ │ │ │ noise_pred=self.unet( │
│ │
│ /mecomics-api/Loras.py:61 in load_lora_weights │
│ │
│ 58 │ │ │ │ curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up, weigh │
│ 59 │ else: │
│ 60 │ │ for ckptpath in checkpoint_path: │
│ ❱ 61 │ │ │ pipeline=load_lora_weights(pipeline, ckptpath, multiplier, device, dtype) │
│ 62 │ return pipeline │
│ 63 def unload_lora_weights(pipeline, checkpoint_path, multiplier, device, dtype): │
│ 64 │ LORA_PREFIX_UNET = "lora_unet" │
│ │
│ /mecomics-api/Loras.py:10 in load_lora_weights │
│ │
│ 7 │ # load LoRA weight from .safetensors │
│ 8 │ if isinstance(checkpoint_path, str): │
│ 9 │ │ │
│ ❱ 10 │ │ state_dict = load_file(checkpoint_path, device=device) │
│ 11 │ │ │
│ 12 │ │ updates = defaultdict(dict) │
│ 13 │ │ for key, value in state_dict.items(): │
│ │
│ /home/alexblattnershalom/.local/lib/python3.8/site-packages/safetensors/torch.py:99 in load_file │
│ │
│ 96 │ ``` │
│ 97 │ """ │
│ 98 │ result = {} │
│ ❱ 99 │ with safe_open(filename, framework="pt", device=device) as f: │
│ 100 │ │ for k in f.keys(): │
│ 101 │ │ │ result[k] = f.get_tensor(k) │
│ 102 │ return result │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
SafetensorError: Error while deserializing header: HeaderTooLarge

I am unsure how to proceed....

[pcuenca]

Hmm, we've seen that error in the past when attempting to load as safetensors files that are not really in safetensors format: #2448, #2551.

Could that be what's happening in your case?

Otherwise, I'd recommend you isolate a short standalone failing code snippet that we can use to reproduce the issue, if possible.

[alexblattner]

hello @pcuenca , thanks for doing your best to help. I decided going for an implementation that uses this: https://gist.github.com/takuma104/e38d683d72b1e448b8d9b3835f7cfa44 instead