flow-matching-cifar10

Use the following code:

from diffusers.models.unets import UNet2DModel
from diffusers.models.embeddings import GaussianFourierProjection

class MyUNet2DModel(UNet2DModel):
    def __init__(self, *args, **kwargs):
        block_out_channels = (128, 256, 256, 256)
        super(MyUNet2DModel, self).__init__(
            in_channels=3, out_channels=3,
            time_embedding_type="fourier",
            down_block_types=("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
            up_block_types=("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
            act_fn="silu",
            block_out_channels=block_out_channels,
            layers_per_block=2, 
            norm_num_groups=32, 
            norm_eps=1e-6,
        )
        self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16,
                                  set_W_to_weight=False, log=False)  # default log=True
        # this cause inf https://github.com/huggingface/diffusers/blob/a536e775fb95daf57abf02dc401e02701591bf69/src/diffusers/models/unets/unet_2d.py#L341
        self.config.time_embedding_type = "my_fourier"  # use a weird name to avoid inf

model = MyUNet2DModel()
model.from_pretrained("Dinghuai/flow-matching-cifar10")