VidMuse_CVPR

model.py

@@ -0,0 +1,17 @@
+# model.py
+from transformers import PreTrainedModel
+import torch
+from audiocraft.models import VidMuse
+from einops import rearrange
+
+class VidMuseModel(PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = VidMuse.get_pretrained(config.model_path)  # 加载你已有的预训练模型
+
+    def forward(self, video_input, **gen_kwargs):
+        # 获取视频的本地帧和全局帧
+        local_video_tensor, global_video_tensor = video_input
+        # 使用 VidMuse 生成音频
+        outputs = self.model.generate([local_video_tensor, global_video_tensor], **gen_kwargs)
+        return outputs

modeling_vidmuse.py

@@ -4,7 +4,7 @@ import torch
 import os
 from huggingface_hub import hf_hub_download
 import os
-
+from huggingface_hub import snapshot_download
 
 # 注册自定义配置和模型（关键步骤！）
 class VidMuseConfig(PretrainedConfig):
@@ -26,24 +26,25 @@ class VidMuseModel(PreTrainedModel):
         )
                     
         self.compression_model = self._load_submodel(config.compression_model)
-    
+        # import pdb; pdb.set_trace()
 
     def _load_submodel(self, relative_path):
         full_path = os.path.join(self.hub_cache_dir, relative_path)
         return torch.load(full_path)
 
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-        # 主动下载附加文件
-        hf_hub_download(
-            repo_id=pretrained_model_name_or_path,
-            filename="compression_state_dict.bin",
-            force_download=True,
-            cache_dir=kwargs.get("cache_dir", None)
-        )
+    # @classmethod
+    # def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+    #     # 主动下载附加文件
+    #     hf_hub_download(
+    #         repo_id=pretrained_model_name_or_path,
+    #         filename="compression_state_dict.bin",
+    #         force_download=True,
+    #         cache_dir=kwargs.get("cache_dir", None)
+    #     )
+        
+    #     # 继续正常加载流程
         
-        # 继续正常加载流程
-        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+    #     return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
     
 # 注册到Auto框架（必须放在类定义之后！）
 AutoConfig.register("vidmuse", VidMuseConfig)

state_dict.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ef70e83c661434c931e6147a35402556bd79f6d0d3d8527205f5ce1ccd26262
+size 7872328846