We're already doing this here no?

Perhaps, guarding this code block with if not is_custom_diffusion would be more helpful?

I hope I give a clear response to your question. :)

Oh thanks so much. In that case, could we maybe simply move

out of

?

I have tried this solution, but it failed. The reason is

It will cause the bug "not same device" when you use GPU to train. (As the self is not been moved to GPU)

In addition I think

can not be simply moved out of

since when PEFT is uninstalled, it should between

and

Hope to help:)

I see. I understand now.

I think the existing

needs to be deleted no?

Then it feels redundant. Because why would we want to add add the same piece of code twice? 😱

Let's make this question clear.

Our target is to let the weights of cross-attention in custom diffusion load successfully whether PEFT is installed since

• PEFT is commonly installed for Lora
• but people who want to use custom_diffusion only don't need PEFT

In my opinion, we need to keep both pieces of code for

• PEFT offload UNet before

so we don't need to offload UNet when PEFT is installed.

• If we delete it. This code may result in unsafety when PEFT is not installed since UNet is not offloaded

Or maybe we can change this piece of code to:

        if not USE_PEFT_BACKEND:
            if _pipeline is not None:
                for _, component in _pipeline.components.items():
                    if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):
                        is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload)
                        is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook)

                        logger.info(
                            "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
                        )
                        remove_hook_from_module(component, recurse=is_sequential_cpu_offload)

        # only custom diffusion needs to set attn processors
        if is_custom_diffusion:
            #Load custom diffusion cross attention weight with PEFT installed in environment
            self.set_attn_processor(attn_processors)

        # set lora layers
        for target_module, lora_layer in lora_layers_list:
            target_module.set_lora_layer(lora_layer)

        self.to(dtype=self.dtype, device=self.device)
        
        if not USE_PEFT_BACKEND:
            # Offload back.
            if is_model_cpu_offload:
                _pipeline.enable_model_cpu_offload()
            elif is_sequential_cpu_offload:
                _pipeline.enable_sequential_cpu_offload()
            # Unsafe code />

Can't we leverage these two variables more systematically to process the logic you mentioned? is_custom_diffusion and USE_PEFT_BACKEND. I am uncomfortable keeping the same exact line of code like this as it makes the code confusing for us, maintainers.

Your proposal looks good to me.

Glad to see that, I will update my PR. It's nice to discuss with you, your idea about code inspires me a lot.