update unet_2d.py and AttnUpBlock2D

[QishengL]

I tried to generate some images based on unet_2d and got the same error in #255. #505 fixed this problem.
As #505, unet_2d_conditon.py and CrossAttnUpBlock2D can work on images that are not multiples of 64, but unet_2d.py and AttnUpBlock2D are not updated.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[patrickvonplaten]

Hey @QishengL,

Thanks for the PR, could you try to add a test as well here? :-)

[QishengL]

from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
import torch

#model_test
model=UNet2DModel(
        sample_size=(80,200),#not multiples of 64
        in_channels=1,
        out_channels=1,
        layers_per_block=2,
        block_out_channels=(128, 128, 256, 256),
        down_block_types=(
            "DownBlock2D",
            "DownBlock2D",
            "AttnDownBlock2D",
            "DownBlock2D",    
        ),
        up_block_types=(
            "UpBlock2D",
            "AttnUpBlock2D",
            "UpBlock2D",
            "UpBlock2D",

        ),
    norm_num_groups= 2
) 
random_image=torch.rand(2,1,80,200)
rand_timesteps=torch.tensor([123, 234])
res=model(random_image,rand_timesteps).sample
print(res.shape) #torch.Size([2, 1, 80, 200])

#pipeline test 
noise_scheduler = DDPMScheduler(
            num_train_timesteps=10,
            beta_schedule="linear",
            prediction_type="epsilon",
        )

pipeline = DDPMPipeline(
                    unet=model,
                    scheduler=noise_scheduler,
            )
images = pipeline(num_inference_steps= 10).images
print(images[0].size) #(200, 80)

[patrickvonplaten]

Hey @QishengL,

Sorry I think I wasn't super clear - I was wondering if you could add a test here:

diffusers/tests/models/test_models_unet_2d.py

Line 32 in 42d9501

class Unet2DModelTests(ModelTesterMixin, unittest.TestCase):

that checks that the unet can process an input that is just a multiple of 8

[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.