dreambooth lora: text encoder monkey-patching not working

[rvorias]

Describe the bug

In the dreambooth lora script, the lora layers are monkey patched into the text encoder as the text encoder comes from transformers. However, this has some undesirable results:

This is before any weight is updated. Vertically, the images share the same exact seed.

I suspect these lines of code:

diffusers/src/diffusers/loaders.py

Lines 949 to 950 in 29b1325

	def new_forward(x):
	return old_forward(x) + lora_layer(x)

Commenting out the lora term produces the same distortions. This means that wrapping the original ref to the forward function and injecting it via monkey patching produces extra undesirable results.

My current guess is that this is some pass-by-ref issue.

Reproduction

train_text_encoder: True and check images at epoch=0 before training.

Logs

No response

System Info

Python 3.8.10

accelerate==0.19.0
diffusers @ git+https://github.com/huggingface/diffusers.git@7a32b6beeb0cfdefed645253dce23d9b0a78597f
transformers==4.29.1

[rvorias]

This seems to alleviate the ref issue

def create_new_forward(old_forward):
    def new_forward(x):
        return old_forward(x)
    return new_forward

# Monkey-patch.
module.forward = create_new_forward(old_forward)

So finally:

...
old_forward = module.forward

def create_new_forward(old_forward):
    def new_forward(x):
        return old_forward(x) + lora_layer(x)
    return new_forward

# Monkey-patch.
module.forward = create_new_forward(old_forward)

[patrickvonplaten]

cc @sayakpaul here

[Mr-Philo]

Same problems here! Adding --train_text_encoder will lead to broken text encoder: the generated images won't follow the input text:

When training after 2500steps, the model will fit the training images but still ignores the input prompt, llike here "the sunflowers":

When disable --train_text_encoder, everything goes smoothly.

Experiment Setting

transformers==4.28.1
accelerate==0.19.0
diffusers=0.17.0.dev0 (latest)

[sayakpaul]

Being worked in #3437. Thanks for your patience.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[patrickvonplaten]

Think we can close this as it's solved, please leave a message if it's not solved for you