ZichengDuan / TheChosenOne

Unofficial implementation of the paper "The Chosen One: Consistent Characters in Text-to-Image Diffusion Models"
https://arxiv.org/abs/2311.10093
242 stars 22 forks source link

text_encoder_lora_state_dict not found #2

Closed zhanchao019 closed 10 months ago

zhanchao019 commented 10 months ago

Facing an env problem: ImportError: cannot import name 'text_encoder_lora_state_dict' from 'diffusers.models.lora'

May I ask which diffusers commit u use? thx

johndpope commented 10 months ago

make sure you have diffusers "0.24.0.dev0"

gh repo clone huggingface/diffusers 
git checkout 2a7f43a
python setup.py install

paste this inside the sdxl...py file


# TODO: This function should be removed once training scripts are rewritten in PEFT
def text_encoder_lora_state_dict(text_encoder):
    state_dict = {}

    def text_encoder_attn_modules(text_encoder):
        from transformers import CLIPTextModel, CLIPTextModelWithProjection

        attn_modules = []

        if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)):
            for i, layer in enumerate(text_encoder.text_model.encoder.layers):
                name = f"text_model.encoder.layers.{i}.self_attn"
                mod = layer.self_attn
                attn_modules.append((name, mod))

        return attn_modules

    for name, module in text_encoder_attn_modules(text_encoder):
        for k, v in module.q_proj.lora_linear_layer.state_dict().items():
            state_dict[f"{name}.q_proj.lora_linear_layer.{k}"] = v

        for k, v in module.k_proj.lora_linear_layer.state_dict().items():
            state_dict[f"{name}.k_proj.lora_linear_layer.{k}"] = v

        for k, v in module.v_proj.lora_linear_layer.state_dict().items():
            state_dict[f"{name}.v_proj.lora_linear_layer.{k}"] = v

        for k, v in module.out_proj.lora_linear_layer.state_dict().items():
            state_dict[f"{name}.out_proj.lora_linear_layer.{k}"] = v

    return state_dict