tentative fix

app.py

@@ -35,16 +35,20 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 # Import args_config module first
 import OmniAvatar.utils.args_config
  
-# Create a minimal args object
+# Create and set global args before any other OmniAvatar imports
 class Args:
     def __init__(self):
         self.rank = 0
+        self.world_size = 1
+        self.local_rank = 0
+        self.device = 'cuda:0'
+        self.num_nodes = 1
         self.dtype = 'bf16'
-        self.exp_path = str(MODELS_DIR / "OmniAvatar-1.3B")
-        self.dit_path = str(MODELS_DIR / "Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors")
-        self.text_encoder_path = str(MODELS_DIR / "Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth")
-        self.vae_path = str(MODELS_DIR / "Wan2.1-T2V-1.3B/Wan2.1_VAE.pth")
-        self.wav2vec_path = str(MODELS_DIR / "wav2vec2-base-960h")
+        self.exp_path = str(Path(os.environ.get('MODELS_DIR', 'pretrained_models')) / "OmniAvatar-1.3B")
+        self.dit_path = str(Path(os.environ.get('MODELS_DIR', 'pretrained_models')) / "Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors")
+        self.text_encoder_path = str(Path(os.environ.get('MODELS_DIR', 'pretrained_models')) / "Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth")
+        self.vae_path = str(Path(os.environ.get('MODELS_DIR', 'pretrained_models')) / "Wan2.1-T2V-1.3B/Wan2.1_VAE.pth")
+        self.wav2vec_path = str(Path(os.environ.get('MODELS_DIR', 'pretrained_models')) / "wav2vec2-base-960h")
         self.train_architecture = 'lora'
         self.lora_rank = 128
         self.lora_alpha = 64.0
@@ -69,6 +73,24 @@ class Args:
         self.image_sizes_720 = [[400, 720], [720, 720], [720, 400]]
         self.image_sizes_1280 = [[720, 720], [528, 960], [960, 528], [720, 1280], [1280, 720]]
         self.seq_len = 200
+        self.infer = True
+        self.debug = False
+    
+    def __contains__(self, key):
+        """Support 'in' operator for checking if attribute exists"""
+        return hasattr(self, key)
+    
+    def __iter__(self):
+        """Make the Args object iterable over its attributes"""
+        return iter(self.__dict__)
+    
+    def keys(self):
+        """Return the attribute names"""
+        return self.__dict__.keys()
+    
+    def __getitem__(self, key):
+        """Support dictionary-style access"""
+        return getattr(self, key)
 
 # Set the global args before any other OmniAvatar imports
 OmniAvatar.utils.args_config.args = Args()