This commit is contained in:
lllyasviel
2023-10-14 02:57:41 -07:00
committed by GitHub
parent 6f24c03826
commit ed69bea3e3
11 changed files with 51 additions and 114 deletions
@@ -296,8 +296,7 @@ class UNetModel(nn.Module):
dims=2,
num_classes=None,
use_checkpoint=False,
use_fp16=False,
use_bf16=False,
dtype=th.float32,
num_heads=-1,
num_head_channels=-1,
num_heads_upsample=-1,
@@ -370,8 +369,7 @@ class UNetModel(nn.Module):
self.conv_resample = conv_resample
self.num_classes = num_classes
self.use_checkpoint = use_checkpoint
self.dtype = th.float16 if use_fp16 else th.float32
self.dtype = th.bfloat16 if use_bf16 else self.dtype
self.dtype = dtype
self.num_heads = num_heads
self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample