import torch
from diffusers import UNet2DModel, DPMSolverSDEScheduler


class FrameConditionalUNet2DModel(UNet2DModel):
    def forward(self, x, t, cond):
        net_input = torch.concat((x, cond / 255 * 2 - 1), dim=1).float()
        output = super().forward(net_input, t, return_dict=False)[0]
        return output