learning/annotated_deep_learning_paper_implementations
1
0
Fork 0
mirror of https://github.com/labmlai/annotated_deep_learning_paper_implementations.git synced 2025-08-14 09:31:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #182 from csl122/patch-1

Browse Source 
add two missing param comments for UNetModel class
This commit is contained in:
Varuna Jayasiri
2023-06-30 11:37:56 +01:00
committed by GitHub
parent 2f7a05da1e 084fc89d63
commit 9c288586bc
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
labml_nn/diffusion/stable_diffusion/model/unet.py
View File

@ -49,7 +49,9 @@ class UNetModel(nn.Module):
:param n_res_blocks: number of residual blocks at each level
:param attention_levels: are the levels at which attention should be performed
:param channel_multipliers: are the multiplicative factors for number of channels for each level
:param n_heads: the number of attention heads in the transformers
:param n_heads: is the number of attention heads in the transformers
:param tf_layers: is the number of transformer layers in the transformers
:param d_cond: is the size of the conditional embedding in the transformers
"""
super().__init__()
self.channels = channels