Repository: huggingface/diffusers
File: src/diffusers/loaders/lora_base.py (Python)

Code before:
        logger.info(f"Loading {prefix}.")
        rank = {}
        state_dict = convert_state_dict_to_diffusers(state_dict)

        # convert state dict
        state_dict = convert_state_dict_to_peft(state_dict)

        for name, _ in text_encoder.named_modules():
            if name.endswith((".q_proj", ".k_proj", ".v_proj", ".out_proj", ".fc1", ".fc2")):
                rank_key = f"{name}.lora_B.weight"
                if rank_key in state_dict:
                    rank[rank_key] = state_dict[rank_key].shape[1]

        if network_alphas is not None:
            alpha_keys = [k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix]
            network_alphas = {k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys}

        # create `LoraConfig`
        lora_config = _create_lora_config(state_dict, network_alphas, metadata, rank, is_unet=False)

        # adapter_name
        if adapter_name is None:
            adapter_name = get_adapter_name(text_encoder)

        # <Unsafe code
        with _lora_loading_context(_pipeline):
            # inject LoRA layers and load the state dict
            # in transformers we automatically check whether the adapter name is already in use or not
            text_encoder.load_adapter(
                adapter_name=adapter_name,
                adapter_state_dict=state_dict,
                peft_config=lora_config,
                **peft_kwargs,
            )

            # scale LoRA layers with `lora_scale`
            scale_lora_layers(text_encoder, weight=lora_scale)

            text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)
        # Unsafe code />

    _maybe_warn_if_no_keys_found(state_dict, prefix, model_class_name=text_encoder.__class__.__name__)


def _func_optionally_disable_offloading(_pipeline):
    is_model_cpu_offload = False
    is_sequential_cpu_offload = False

    if _pipeline is not None and _pipeline.hf_device_map is None:
        for _, component in _pipeline.components.items():
            if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):

Review comment (bug):
I'm not sure if a context manager is helping here, what was wrong with https://github.com/huggingface/diffusers/pull/11719/files#diff-005343920b6fa6c3f4c1f639ee23f180dfa424d6008e487fdc9962a82f498d15L448-L452?

Code after:
        logger.info(f"Loading {prefix}.")
        rank = {}
        state_dict = convert_state_dict_to_diffusers(state_dict)

        # convert state dict
        state_dict = convert_state_dict_to_peft(state_dict)

        for name, _ in text_encoder.named_modules():
            if name.endswith((".q_proj", ".k_proj", ".v_proj", ".out_proj", ".fc1", ".fc2")):
                rank_key = f"{name}.lora_B.weight"
                if rank_key in state_dict:
                    rank[rank_key] = state_dict[rank_key].shape[1]

        if network_alphas is not None:
            alpha_keys = [k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix]
            network_alphas = {k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys}

        # create `LoraConfig`
        lora_config = _create_lora_config(state_dict, network_alphas, metadata, rank, is_unet=False)

        # adapter_name
        if adapter_name is None:
            adapter_name = get_adapter_name(text_encoder)

        # <Unsafe code
        is_model_cpu_offload, is_sequential_cpu_offload = _func_optionally_disable_offloading(_pipeline)
        # inject LoRA layers and load the state dict
        # in transformers we automatically check whether the adapter name is already in use or not
        text_encoder.load_adapter(
            adapter_name=adapter_name,
            adapter_state_dict=state_dict,
            peft_config=lora_config,
            **peft_kwargs,
        )

        # scale LoRA layers with `lora_scale`
        scale_lora_layers(text_encoder, weight=lora_scale)
        text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)

        # Offload back.
        if is_model_cpu_offload:
            _pipeline.enable_model_cpu_offload()
        elif is_sequential_cpu_offload:
            _pipeline.enable_sequential_cpu_offload()
        # Unsafe code />

    if prefix is not None and not state_dict:
        model_class_name = text_encoder.__class__.__name__
        logger.warning(
            f"No LoRA keys associated to {model_class_name} found with the {prefix=}. "
            "This is safe to ignore if LoRA state dict didn't originally have any "