Fix UNET and VAE implementations for new diffusers version (#4663)

This PR updates the DepSpeed UNET and VAE implementations to support `diffusers>=0.23.0`.
deepspeedai · Nov 10, 2023 · a361bac · a361bac
1 parent 4388a60
commit a361bac
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/deepspeed/model_implementations/diffusers/unet.py b/deepspeed/model_implementations/diffusers/unet.py
@@ -62,7 +62,13 @@ def _create_cuda_graph(self, *inputs, **kwargs):
 
         self.cuda_graph_created = True
 
-    def _forward(self, sample, timestamp, encoder_hidden_states, return_dict=True, cross_attention_kwargs=None):
+    def _forward(self,
+                 sample,
+                 timestamp,
+                 encoder_hidden_states,
+                 return_dict=True,
+                 cross_attention_kwargs=None,
+                 timestep_cond=None):
         if cross_attention_kwargs:
             return self.unet(sample,
                              timestamp,

diff --git a/deepspeed/model_implementations/diffusers/vae.py b/deepspeed/model_implementations/diffusers/vae.py
@@ -30,7 +30,7 @@ def _graph_replay_decoder(self, *inputs, **kwargs):
         self._decoder_cuda_graph.replay()
         return self.static_decoder_output
 
-    def _decode(self, x, return_dict=True):
+    def _decode(self, x, return_dict=True, generator=None):
         return self.vae.decode(x, return_dict=return_dict)
 
     def _create_cuda_graph_decoder(self, *inputs, **kwargs):