import fixes.

xora/models/autoencoders/video_autoencoder.py

@@ -11,7 +11,7 @@ from torch.nn import functional
 
 from diffusers.utils import logging
 
-from txt2img.models.layers.nn import Identity
+from xora.utils.torch_utils import Identity
 from xora.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd
 from xora.models.autoencoders.pixel_norm import PixelNorm
 from xora.models.autoencoders.vae import AutoencoderKLWrapper

xora/schedulers/rf.py

@@ -9,7 +9,7 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from diffusers.utils import BaseOutput
 from torch import Tensor
 
-from txt2img.common.torch_utils import append_dims
+from xora.utils.torch_utils import append_dims
 
 
 def simple_diffusion_resolution_dependent_timestep_shift(

xora/utils/torch_utils.py

@@ -1,4 +1,5 @@
 import torch
+from torch import nn
 
 def append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor:
     """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
@@ -8,3 +9,13 @@ def append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor:
     elif dims_to_append == 0:
         return x
     return x[(...,) + (None,) * dims_to_append]
+
+class Identity(nn.Module):
+    """A placeholder identity operator that is argument-insensitive."""
+
+    def __init__(self, *args, **kwargs) -> None:  # pylint: disable=unused-argument
+        super().__init__()
+
+    # pylint: disable=unused-argument
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        return x