Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -71,6 +71,7 @@ from finetune.datasets.utils import (
 from torch.utils.data import Dataset
 from torchvision import transforms
 
+from safetensors.torch import load_file
 
 import pdb
 sys.path.append(os.path.abspath(os.path.join(sys.path[-1], 'finetune'))) # for camera flow generator
@@ -143,7 +144,9 @@ def load_cogvideox_flovd_OMSM_lora_pipeline(omsm_path, backbone_path, transforme
     # 1) Load Lora weight
     transformer.add_adapter(transformer_lora_config)
 
-    lora_state_dict = FloVDOMSMCogVideoXImageToVideoPipeline.lora_state_dict(omsm_path)
+    lora_path = os.path.join(omsm_path, "pytorch_lora_weights.safetensors")
+    lora_state_dict = load_file(lora_path)
+
     transformer_state_dict = {
         f'{k.replace("transformer.", "")}': v
         for k, v in lora_state_dict.items()