fix an issue with karras unet when no class labels, and make sure it is compatible with GaussianDiffusion wrapper

Author: lucidrains

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -489,7 +489,7 @@ def __init__(
     ):
         super().__init__()
         assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
-        assert not model.random_or_learned_sinusoidal_cond
+        assert not hasattr(model, 'random_or_learned_sinusoidal_cond') or not model.random_or_learned_sinusoidal_cond
 
         self.model = model
 

denoising_diffusion_pytorch/karras_unet.py

@@ -1,3 +1,7 @@
+"""
+the magnitude-preserving unet proposed in https://arxiv.org/abs/2312.02696 by Karras et al.
+"""
+
 import math
 from math import sqrt, ceil
 from functools import partial
@@ -456,8 +460,9 @@ def __init__(
         self.needs_class_labels = exists(num_classes)
         self.num_classes = num_classes
 
-        self.to_class_emb = Linear(num_classes, 4 * dim)
-        self.add_class_emb = MPAdd(t = mp_add_emb_t)
+        if self.needs_class_labels:
+            self.to_class_emb = Linear(num_classes, 4 * dim)
+            self.add_class_emb = MPAdd(t = mp_add_emb_t)
 
         # final embedding activations
 
@@ -537,6 +542,8 @@ def __init__(
             Decoder(curr_dim, curr_dim, has_attn = mid_has_attn, **block_kwargs),
         ])
 
+        self.out_dim = channels
+
     @property
     def downsample_factor(self):
         return 2 ** self.num_downsamples

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.10.1'
+__version__ = '1.10.2'