Train nyu dataset suffer error about datasetdm

[hanyuxuan123]

Hello, I try the code and face the same problem
I use command:
python tools/splicing_augmentation_Cityscapes.py --input_root "CityScape" --out_root "Augmentation" --thread_num 7

In the File "/workspace/my_code/DatasetDM/model/unet.py", line 153, in forward
reshape_h=h.reshape(int(h.size()[0]/2),int(h.size()[1]*2),h.size()[2],h.size()[3])
IndexError:tuple index out of range

I debuged the code and found that

the upsample code is
sample, res_samples= upsample_block(
hidden_states=sample,
temb=emb,
res_hidden_states_tuple=res_samples,
encoder_hidden_states=encoder_hidden_states,
)
However in the diffusers the diffusers/src/diffusers/models/unet_2D_blocks.py class CrossAttnUpBlock2D and class UpBlock2D the 'def forward' only have one return variable called hidden_states in type of 'tensors' [2,1280,8,8]. the function above uses two variable samples and res_samples, it is separated to [1280,8,8] and [1280,8,8]. However, in the previous downsample part the class return a tensor(hidden states) and tuple(output_samples). In the for loop the [1280,8,8] was separated to [8,8], so the tuple index out of range.

I recommend to use another code:
in the unet.py
in the upsample part new code:
for upsample_block in self.up_blocks:
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
if hasattr(upsample_block, "attentions") and upsample_block.attentions is not None:
sample= upsample_block(
hidden_states=sample,
temb=emb,
res_hidden_states_tuple=res_samples,
encoder_hidden_states=encoder_hidden_states,
)
else:
sample = upsample_block(hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples)

        if flag_time:
            # reshape_h=sample.reshape(int(h.size()[0]/2),int(h.size()[1]*2),h.size()[2],h.size()[3])
            print(sample.shape)
            reshape_h=sample.reshape(int(sample.size()[0]/2),int(sample.size()[1]*2),sample.size()[2],sample.size()[3])

reshape_h=sample.reshape(int(sample.size()[0]/2),int(sample.size()[1]*2),sample.size()[2],sample.size()[3])

            if reshape_h.size()[2]==8:
                all_feature_dic["low"].append(reshape_h)
            elif reshape_h.size()[2]==16:
                all_feature_dic["mid"].append(reshape_h)
            elif reshape_h.size()[2]==32:
                all_feature_dic["high"].append(reshape_h)
            elif reshape_h.size()[2]==64:
                all_feature_dic["highest"].append(reshape_h)

[hanyuxuan123]

    for upsample_block in self.up_blocks:
        res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
        if hasattr(upsample_block, "attentions") and upsample_block.attentions is not None:
            sample= upsample_block(
                hidden_states=sample,
                temb=emb,
                res_hidden_states_tuple=res_samples,
                encoder_hidden_states=encoder_hidden_states,
            )
        else:
            sample = upsample_block(hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples)
        
        if flag_time:
            # reshape_h=sample.reshape(int(h.size()[0]/2),int(h.size()[1]*2),h.size()[2],h.size()[3])
            print(sample.shape)
            reshape_h=sample.reshape(int(sample.size()[0]/2),int(sample.size()[1]*2),sample.size()[2],sample.size()[3])
            if reshape_h.size()[2]==8:
                all_feature_dic["low"].append(reshape_h)
            elif reshape_h.size()[2]==16:
                all_feature_dic["mid"].append(reshape_h)
            elif reshape_h.size()[2]==32:
                all_feature_dic["high"].append(reshape_h)
            elif reshape_h.size()[2]==64:
                all_feature_dic["highest"].append(reshape_h)