Update import for changes to latest diffusers (deepspeedai#5065)

HabanaAI · Feb 2, 2024 · a049370 · a049370
1 parent 567f97b
commit a049370
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepspeed/module_inject/containers/unet.py b/deepspeed/module_inject/containers/unet.py
@@ -17,6 +17,8 @@ def __init__(self):
         try:
             import diffusers
             self._orig_layer_class = diffusers.models.unet_2d_condition.UNet2DConditionModel
+        except AttributeError:
+            self._orig_layer_class = diffusers.models.unets.unet_2d_condition.UNet2DConditionModel
         except ImportError:
             self._orig_layer_class = None