`UNet2DOutput` becomes nested when wrapped with `accelerate`

[anton-l]

Describe the bug

When either UNet2DModel or UNet2DConditionModel are prepared with accelerate, their outputs become nested, i.e. to get the sample you have to do outputs.sample['sample'] instead of just outputs.sample.

However, it works as expected with return_dict=False.

Linked to #499 #434

cc @patrickvonplaten @patil-suraj @sgugger

Reproduction

• Regular outputs:

from accelerate import Accelerator
from diffusers import UNet2DModel


model = UNet2DModel(
    sample_size=16,
    in_channels=3,
    out_channels=3,
    layers_per_block=1,
    block_out_channels=(128,),
    down_block_types=("DownBlock2D",), 
    up_block_types=("UpBlock2D",),
)
model = model.cuda()

x = torch.randn((1, 3, 16, 16)).to(model.device)
t = torch.tensor([0]).to(model.device)
print(model(x, t))

UNet2DOutput(sample=tensor([[[[ 6.6783e-02,  1.1064e-01, -3.1808e-01, -2.7287e-01,  7.0199e-02,
           -1.7103e-01,  4.0218e-01,  1.7775e-01,  6.2583e-01, -1.4165e-01,
            9.0453e-02, -8.6686e-02, -3.2533e-01,  3.4424e-02,  4.2162e-02,
            8.4619e-02],
...

• Outputs after accelerator.prepare:

accelerator = Accelerator()
model = accelerator.prepare(model)

print(model(x, t))

UNet2DOutput(sample={'sample': tensor([[[[ 6.6467e-02,  1.1072e-01, -3.1836e-01, -2.7295e-01,  6.9458e-02,
           -1.7090e-01,  4.0234e-01,  1.7761e-01,  6.2598e-01, -1.4160e-01,
            9.0271e-02, -8.6182e-02, -3.2568e-01,  3.4332e-02,  4.2419e-02,
            8.4656e-02],

Logs

No response

System Info

diffusers: 0.3.0
accelerate: 0.12.0
torch: 1.12.1+cu113

single GPU, colab

[sgugger]

Copying my internal answer here for posterity :-)

Basically a dictionary in Python can be instantiated form kwargs:

DictClass(a=1, b=2, c=3)

or a dict

DictClass({"a": 1, "b":2, "c":3})

More genertlly any kind of iterator that produces key value pairs. Some classes only support the second kind, so in Accelerate when we treat the output of the models and have to reconstruct them, we use this way (specifically type(output)(processed_dict)

Now the BaseOutput class does not support the second assignment:

output = UNet2DOutput({"sample": torch.randn(3, 4)})

will return a level of nestedness too much, when it should just do the same thing as

output = UNet2DOutput(sample = torch.randn(3, 4))

This is how we handle it in Transformers.