|
--- |
|
library_name: diffusers |
|
tags: |
|
- diffusion |
|
- flow-matching |
|
- diffusers |
|
--- |
|
|
|
Use the following code: |
|
```python |
|
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") |
|
``` |