Function create_custom_diffusion

src/diffusers_training_sdxl.py:267–297

(unet, freeze_model)

def create_custom_diffusion(unet, freeze_model):
    for name, params in unet.named_parameters():
        if freeze_model == 'crossattn':
            if 'attn2' in name:
                params.requires_grad = True
                print(name)
            else:
                params.requires_grad = False
        elif freeze_model == "crossattn_kv":
            if 'attn2.to_k' in name or 'attn2.to_v' in name:
                params.requires_grad = True
                print(name)
            else:
                params.requires_grad = False
        else:
            raise ValueError(
                    "freeze_model argument only supports crossattn_kv or crossattn"
                )

    # change attn class
    def change_attn(unet):
        for layer in unet.children():
            if type(layer) is Attention:
                bound_method = set_use_memory_efficient_attention_xformers.__get__(layer, layer.__class__)
                setattr(layer, 'set_use_memory_efficient_attention_xformers', bound_method)
            else:
                change_attn(layer)

    change_attn(unet)
    unet.set_attn_processor(CustomDiffusionAttnProcessor())
    return unet


# Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt