Skip to content

GitLab

Unverified Commit 88b015dc authored by Xinyuan Zhao's avatar Xinyuan Zhao Committed by GitHub
Browse files

Make `time_embed_dim` of `UNet2DModel` changeable (#10262)

parent 63cdf9c0
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
src/diffusers/models/unets/unet_2d.py
View file @ 88b015dc
...@@ -97,6 +97,7 @@ class UNet2DModel(ModelMixin, ConfigMixin): ...@@ -97,6 +97,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
out_channels: int = 3, out_channels: int = 3,
center_input_sample: bool = False, center_input_sample: bool = False,
time_embedding_type: str = "positional", time_embedding_type: str = "positional",
time_embedding_dim: Optional[int] = None,
freq_shift: int = 0, freq_shift: int = 0,
flip_sin_to_cos: bool = True, flip_sin_to_cos: bool = True,
down_block_types: Tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"), down_block_types: Tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
...@@ -122,7 +123,7 @@ class UNet2DModel(ModelMixin, ConfigMixin): ...@@ -122,7 +123,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
super().__init__() super().__init__()
self.sample_size = sample_size self.sample_size = sample_size
time_embed_dim = block_out_channels[0] * 4 time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
# Check inputs # Check inputs
if len(down_block_types) != len(up_block_types): if len(down_block_types) != len(up_block_types):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment