make einops optional for RL

1a099e5e · anton-l · b09b152f · 1a099e5e
Commit 1a099e5e authored Jun 21, 2022 by anton-l
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

src/diffusers/models/unet_rl.py src/diffusers/models/unet_rl.py +6 -2

No files found.
--- a/src/diffusers/models/unet_rl.py
+++ b/src/diffusers/models/unet_rl.py
@@ -5,8 +5,12 @@ import math
 import torch
 import torch.nn as nn
-import einops
+try:
-from einops.layers.torch import Rearrange
+    import einops
+    from einops.layers.torch import Rearrange
+except:
+    print("Einops is not installed")
+    pass
 from ..configuration_utils import ConfigMixin