Update models/__init__.py

models/__init__.py

@@ -2,8 +2,6 @@ import os
 import sys
 sys.path.append(os.path.split(sys.path[0])[0])
 
-from .dit import DiT_models
-from .uvit import UViT_models
 from .unet import UNet3DConditionModel
 from torch.optim.lr_scheduler import LambdaLR
 
@@ -28,25 +26,7 @@ def get_lr_scheduler(optimizer, name, **kwargs):
     
 def get_models(args):
     
-    if 'DiT' in args.model:
-        return DiT_models[args.model](
-                input_size=args.latent_size,
-                num_classes=args.num_classes,
-                class_guided=args.class_guided,
-                num_frames=args.num_frames,
-                use_lora=args.use_lora,
-                attention_mode=args.attention_mode
-            )
-    elif 'UViT' in args.model:
-        return UViT_models[args.model](
-                input_size=args.latent_size,
-                num_classes=args.num_classes,
-                class_guided=args.class_guided,
-                num_frames=args.num_frames,
-                use_lora=args.use_lora,
-                attention_mode=args.attention_mode
-            )
-    elif 'TAV' in args.model:
+    if 'TAV' in args.model:
         pretrained_model_path = args.pretrained_model_path
         return UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", use_concat=args.use_mask)
     else: