fix

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -13,7 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from typing import Optional, Tuple, Any, Dict, Union
+from typing import Optional, Tuple, Any, Dict
 import jax
 import jax.numpy as jnp
 from flax import nnx
@@ -23,8 +23,6 @@
 from maxdiffusion.models.attention_flax import NNXSimpleFeedForward
 from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection
 from maxdiffusion.models.gradient_checkpoint import GradientCheckpointType
-from maxdiffusion.common_types import BlockSizes
-from maxdiffusion.configuration_utils import ConfigMixin, register_to_config
 from maxdiffusion.configuration_utils import ConfigMixin, register_to_config