Merge pull request #36 from LightricksResearch/bugfix/check_timesteps_eligibility

xora/models/autoencoders/vae.py

@@ -1,6 +1,7 @@
 from typing import Optional, Union
 
 import torch
+import inspect
 import math
 import torch.nn as nn
 from diffusers import ConfigMixin, ModelMixin
@@ -60,6 +61,8 @@ class AutoencoderKLWrapper(ModelMixin, ConfigMixin):
         self.dims = dims
         self.z_sample_size = 1
 
+        self.decoder_params = inspect.signature(self.decoder.forward).parameters
+
         # only relevant if vae tiling is enabled
         self.set_tiling_params(sample_size=sample_size, overlap_factor=0.25)
 
@@ -257,7 +260,10 @@ class AutoencoderKLWrapper(ModelMixin, ConfigMixin):
         timesteps: Optional[torch.Tensor] = None,
     ) -> Union[DecoderOutput, torch.FloatTensor]:
         z = self.post_quant_conv(z)
-        dec = self.decoder(z, target_shape=target_shape, timesteps=timesteps)
+        if "timesteps" in self.decoder_params:
+            dec = self.decoder(z, target_shape=target_shape, timesteps=timesteps)
+        else:
+            dec = self.decoder(z, target_shape=target_shape)
         return dec
 
     def decode(