Possible bug in ControlNet code

[treasan]

Hey,

I have a question regarding the PixArt ControlNet code.

The paper suggests to add the input directly to the conditioning signal in the trainable network copy. However in the code, you first forward the input through the first frozen block and only then add it to the conditioning.

The first trainable copy therefore does not receive the input as expected, but the already modified and slightly encoded input from the first frozen layer. This seems unintuitive. Is it possible that this is wrong behaviour or did I overlook something?

[sanshibayuan]

Did you fix this? I also found the ControlNet result is not right compared to the paper, is it possible causing the problem? Did you test it ?

[lawrence-cj]

This issue may cause problem when the input is image, not HED. We have fixed the bug in the ControlNet in PixArt-Sigma. Will released soon. Stay tune!

https://github.com/PixArt-alpha/PixArt-sigma

[Feynman1999]

Will inputting an image (not a hed) cause a bug? If I modify the code logic myself and add input to the condition at the beginning, is that okay?

[lawrence-cj]

If your input is an image, the training may be unstable.

If I modify the code logic myself and add input to the condition at the beginning, is that okay?

Yes. It definitely will help with the task with image inputs.

[Feynman1999]

Hey,

I have a question regarding the PixArt ControlNet code.

The paper suggests to add the input directly to the conditioning signal in the trainable network copy. However in the code, you first forward the input through the first frozen block and only then add it to the conditioning.

The first trainable copy therefore does not receive the input as expected, but the already modified and slightly encoded input from the first frozen layer. This seems unintuitive. Is it possible that this is wrong behaviour or did I overlook something?

I have found the same problem as you, that is, there is a slight difference between the process illustrated in the paper and the actual code. What I'm thinking is that since the copied block is trainable, the difference may not be significant?

[Feynman1999]

I am currently trying to use Pixart alpha to train ControlNet on my own dataset, with a training resolution of 1024 and a task of image restoration (condition is image). If there is any progress, I will synchronize it here

[lawrence-cj]

You can try to modify the bug in your code base. Our new version controlnet fixed this bug.

[brycegoh]

What exactly should we edit? Is there a PR we could reference on PixArt-sigma as you mentioned that it is fixed there?

Thanks!

[k-zha14]

Nice observation, thanks bro! @Tomsen1410
For anyone who is confused with the disscussion above, I provide a correction example for reference. @brycegoh CC @lawrence-cj

line 116~122 in PixArt-alpha/diffusion/model/nets/pixart_controlnet.py

x_init  = x.detach().clone()    # cache the initial x, fix bug
x = auto_grad_checkpoint(self.base_model.blocks[0], x, y, t0, y_lens, **kwargs)  # (N, T, D) #support grad checkpoint

if c is not None:
    # update c
    for index in range(1, self.copy_blocks_num + 1):
        if index==1:
            c, c_skip = auto_grad_checkpoint(self.controlnet[index - 1], x_init, y, t0, y_lens, c, **kwargs)
        else:
            c, c_skip = auto_grad_checkpoint(self.controlnet[index - 1], x, y, t0, y_lens, c, **kwargs)
        x = auto_grad_checkpoint(self.base_model.blocks[index], x + c_skip, y, t0, y_lens, **kwargs)

[sawyer7246]

Hi, I get another issue when training the controlnet on my own custom dataset. According to the training log, the loss does not decrease and the grad_norm is constantly 0.00.

[Feynman1999]

Hi, I get another issue when training the controlnet on my own custom dataset. According to the training log, the loss does not decrease and the grad_norm is constantly 0.00.

if the grad_norm is constantly 0.0 , I think it's possible that there were bugs in your changes, such as the require_grads are all false, etc

[k-zha14]

Interesting, I work with the modified version successfully. Actually, the x_init is only a value copy. So the 'detach' should have NO influence on the trainable parameters. Hoping for your further feedback!

[sawyer7246]

Hi, I get another issue when training the controlnet on my own custom dataset. According to the training log, the loss does not decrease and the grad_norm is constantly 0.00.

if the grad_norm is constantly 0.0 , I think it's possible that there were bugs in your changes, such as the require_grads are all false, etc

The require_grad attribute of the loss term is True. I also notice a lot of warnings appearing after the weights are loaded, any clues?

WARNING - Missing keys: ['pos_embed', 'x_embedder.proj.weight', 'x_embedder.proj.bias', 't_embedder.mlp.0.weight', 't_embedder.mlp.0.bias', 't_embedder.mlp.2.weight', 't_embedder.mlp.2.bias', 't_block.1.weight', 't_block.1.bias', 'y_embedder.y_embedding', 'y_embedder.y_proj.fc1.weight', 'y_embedder.y_proj.fc1.bias', 'y_embedder.y_proj.fc2.weight', 'y_embedder.y_proj.fc2.bias', 'blocks.0.scale_shift_table', 'blocks.0.attn.qkv.weight', 'blocks.0.attn.qkv.bias', 'blocks.0.attn.proj.weight', 'blocks.0.attn.proj.bias', 'blocks.0.cross_attn.q_linear.weight', 'blocks.0.cross_attn.q_linear.bias', 'blocks.0.cross_attn.kv_linear.weight', 'blocks.0.cross_attn.kv_linear.bias', 'blocks.0.cross_attn.proj.weight', 'blocks.0.cross_attn.proj.bias', 'blocks.0.mlp.fc1.weight', 'blocks.0.mlp.fc1.bias', 'blocks.0.mlp.fc2.weight', 'blocks.0.mlp.fc2.bias', 'blocks.1.scale_shift_table', 'blocks.1.attn.qkv.weight', 'blocks.1.attn.qkv.bias', 'blocks.1.attn.proj.weight', 'blocks.1.attn.proj.bias', 'blocks.1.cross_attn.q_linear.weight', 'blocks.1.cross_attn.q_linear.bias', 'blocks.1.cross_attn.kv_linear.weight', 'blocks.1.cross_attn.kv_linear.bias', 'blocks.1.cross_attn.proj.weight', 'blocks.1.cross_attn.proj.bias', 'blocks.1.mlp.fc1.weight'

[sawyer7246]

more details on this.
The x becomes zero after the final layer.

# ... x has non-zero value.
x = self.final_layer(x, t)
# everything in x become zero.

[sawyer7246]

sorry, it's my fault, I just load a wrong checkpoint ...

load_from = should be a PixArt-XL-2-1024-MS.pth not ControlNet.pth...