UNet3D

Author: hlky

src/diffusers/__init__.py

@@ -198,6 +198,7 @@
             "UNet1DModel",
             "UNet2DConditionModel",
             "UNet2DModel",
+            "UNet3DModel",
             "UNet3DConditionModel",
             "UNetControlNetXSModel",
             "UNetMotionModel",
@@ -761,6 +762,7 @@
             UNet2DConditionModel,
             UNet2DModel,
             UNet3DConditionModel,
+            UNet3DModel,
             UNetControlNetXSModel,
             UNetMotionModel,
             UNetSpatioTemporalConditionModel,

src/diffusers/models/__init__.py

@@ -86,6 +86,7 @@
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
+    _import_structure["unets.unet_3d"] = ["UNet3DModel"]
     _import_structure["unets.unet_3d_condition"] = ["UNet3DConditionModel"]
     _import_structure["unets.unet_i2vgen_xl"] = ["I2VGenXLUNet"]
     _import_structure["unets.unet_kandinsky3"] = ["Kandinsky3UNet"]
@@ -176,6 +177,7 @@
             UNet2DConditionModel,
             UNet2DModel,
             UNet3DConditionModel,
+            UNet3DModel,
             UNetMotionModel,
             UNetSpatioTemporalConditionModel,
             UVit2DModel,

src/diffusers/models/unets/__init__.py

@@ -5,6 +5,7 @@
     from .unet_1d import UNet1DModel
     from .unet_2d import UNet2DModel
     from .unet_2d_condition import UNet2DConditionModel
+    from .unet_3d import UNet3DModel
     from .unet_3d_condition import UNet3DConditionModel
     from .unet_i2vgen_xl import I2VGenXLUNet
     from .unet_kandinsky3 import Kandinsky3UNet