{{ message }}
Repository navigation
Help replicating imagen #1348
Description
Hello, I'm trying to replicate imagen using the diffusers library, so far I've trained the 64x64 unet, next I have to train the 256x256 model.
My question is, what is the best parameters for the second unet?
I am currently using the following configuration, but not yet getting good results
_class_name="UNet2DConditionModel",
act_fn = "silu",
attention_head_dim= 16,
block_out_channels= [
160,
320,
320,
640
],
center_input_sample= False,
cross_attention_dim= 768,
down_block_types= [
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"CrossAttnDownBlock2D"
],
downsample_padding= 1,
flip_sin_to_cos= True,
freq_shift= 0,
in_channels= 6,
layers_per_block= 2,
mid_block_scale_factor= 1,
norm_eps= 1e-05,
norm_num_groups= 32,
out_channels= 3,
sample_size= 256,
up_block_types= [
"CrossAttnUpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D"
]
This is the output of the first unet
this is the output of the second unet
Maybe using the original configuration from the imagen paper would be a good start, but I don't know how the configuration reported by them can be translated into the diffusers framework, any ideas?
Here is the imagen reported config:
architecture = {
"dropout": 0.0,
"feature_pooling_type": "attention",
"use_scale_shift_norm": True,
"blocks": [
{
"channels": 128,
"strides": (2, 2),
"kernel_size": (3, 3),
"num_res_blocks": 2,
},
{
"channels": 256,
"strides": (2, 2),
"kernel_size": (3, 3),
"num_res_blocks": 4,
},
{
"channels": 512,
"strides": (2, 2),
"kernel_size": (3, 3),
"num_res_blocks": 8,
},
{
"channels": 1024,
"strides": (2, 2),
"kernel_size": (3, 3),
"num_res_blocks": 8,
"self_attention": True,
"text_cross_attention": True,
"num_attention_heads": 8
}
]
}
Activity
patrickvonplaten commentedon Nov 21, 2022
Good question! To be honest, I don't really know - does someone from the community maybe know? Also maybe gently pinging @lucidrains here
thedarkzeno commentedon Nov 21, 2022
I think that we need to specify the number of resnet in each block and also the number of attention heads rather than the attention dim. This means that we would have to change the unet code, right?
jenkspt commentedon Nov 21, 2022
@thedarkzeno are you training a super-resolution model for the "256x256" model?
thedarkzeno commentedon Nov 21, 2022
@jenkspt Yes, from 64 to 256 like in the imagen paper
jenkspt commentedon Nov 21, 2022
I guess my question is, are you passing in the low res conditioning image? Because I don't see that as an option for
UNet2DConditionModel. I would expect you would need a subclass like so:(this is untested)
One of the problems with this however is the sample has shape (B, 3, H, W), but the in_channels attribute will be 6
thedarkzeno commentedon Nov 22, 2022
That is exactly what I am doing, I am adding the streched low res image in the extra channels, this way I have 6 channels in the input and 3 in the output. My problem is with the params of the unet like channel size in the blocks, number of attention heads and number of resnets
github-actions commentedon Dec 21, 2022
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.