File size: 1,374 Bytes
363d564 d518c25 363d564 c677162 75936c0 c677162 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
---
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")
``` |