support patch ascend

Dockerfile.ascend

@@ -0,0 +1,9 @@
+FROM git.modelhub.org.cn:9443/enginex-ascend/vllm-ascend:v0.10.0rc1
+
+WORKDIR /workspace
+
+RUN pip install diffusers==0.34.0
+RUN pip install imageio[ffmpeg] einops datasets==3.2.0 simplejson addict open_clip_torch==2.24.0 sortedcontainers modelscope==1.28.2 av==11.0.0 pytorch-lightning
+COPY . /workspace/
+RUN patch /usr/local/python3.11.13/lib/python3.11/site-packages/modelscope/models/multi_modal/video_synthesis/text_to_video_synthesis_model.py patch.ascend/text_to_video_synthesis_model.py.patch
+RUN patch /usr/local/python3.11.13/lib/python3.11/site-packages/modelscope/utils/device.py patch.ascend/device.py.patch

iic.py

@@ -1,26 +1,16 @@
 import os
 import torch
-from functools import wraps
-
-_orig_load = torch.load
-
-@wraps(_orig_load)
-def _load_patch(*args, **kwargs):
-    kwargs.setdefault("weights_only", False)
-    return _orig_load(*args, **kwargs)
-
-torch.load = _load_patch
-
+import patch
 
 from modelscope.pipelines import pipeline
 from modelscope.outputs import OutputKeys
 
+device = "cuda" if torch.cuda.is_available() else "npu"
 model_path = "/mnt/contest_ceph/zhanghao/models/iic/text-to-video-synthesis"
-p = pipeline('text-to-video-synthesis', model_path)
+p = pipeline('text-to-video-synthesis', model_path, device=device)
 
 test_text = {
-        'text': 'A panda eating bamboo on a rock.',
+        'text': 'A panda eating a burger and french fries on a rock.',
     }
-output_video_path = p(test_text, output_video='./output.mp4')[OutputKeys.OUTPUT_VIDEO]
+output_video_path = p(test_text, device=device, output_video='./output.mp4')[OutputKeys.OUTPUT_VIDEO]
 print('output_video_path:', output_video_path)
-

main.py

@@ -8,21 +8,11 @@ import re
 import time
 from datetime import datetime
 from pathlib import Path
-
-import torch
-from functools import wraps
-
-_orig_load = torch.load
-
-@wraps(_orig_load)
-def _load_patch(*args, **kwargs):
-    kwargs.setdefault("weights_only", False)
-    return _orig_load(*args, **kwargs)
-
-torch.load = _load_patch
+import patch
 
 from modelscope.pipelines import pipeline
 from modelscope.outputs import OutputKeys
+import torch
 
 
 def safe_stem(text: str, maxlen: int = 60) -> str:
@@ -58,7 +48,7 @@ def load_prompts(json_path: Path):
 
 
 def build_pipeline(model_path: str, device: str = "cuda", dtype=torch.float16):
-    pipe = pipeline('text-to-video-synthesis', model_path)
+    pipe = pipeline('text-to-video-synthesis', model_path, device=device)
     return pipe
 
 
@@ -95,7 +85,7 @@ def main():
     parser.add_argument("--json", required=True, help="测试文本 JSON 文件路径")
     parser.add_argument("--results", required=True, help="结果 JSON 文件输出路径（*.json）")
     parser.add_argument("--outdir", required=True, help="图片输出目录")
-    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"], help="推理设备")
+    parser.add_argument("--device", default="cuda", help="推理设备")
     parser.add_argument("--dtype", default="fp16", choices=["fp16", "fp32"], help="推理精度")
     args, _ = parser.parse_known_args()
 

patch.ascend/device.py.patch

@@ -0,0 +1,2 @@
+27d26
+<     assert eles[0] in ['cpu', 'cuda', 'gpu'], err_msg

patch.ascend/text_to_video_synthesis_model.py.patch

@@ -0,0 +1,10 @@
+60a61
+>         print(f"kwargs: {kwargs}")
+62a64
+>         print(f"device: {self.device}")
+129c131
+<             layer='penultimate')
+---
+>             layer='penultimate', device=self.device)
+224a227
+>         print(f"self.device: {self.device}")

patch.py

@@ -0,0 +1,43 @@
+import torch
+from functools import wraps
+
+def enable_cuda_to_npu_shim():
+    print("enable_cuda_to_npu_shim")
+    import torch_npu  # 注册 npu
+    # 仅映射常见函数；不要贪多
+    torch.cuda.is_available  = torch.npu.is_available
+    torch.cuda.device_count  = torch.npu.device_count
+    torch.cuda.current_device= torch.npu.current_device
+    torch.cuda.set_device    = torch.npu.set_device
+    torch.cuda.synchronize   = torch.npu.synchronize
+    try:
+        # 若存在空缓存接口
+        torch.cuda.empty_cache = torch.npu.empty_cache   # 某些版本可用
+    except Exception:
+        pass
+    # 设备字符串统一用 npu
+    # 业务里仍建议 model.to("npu:0") 显式写清
+try:
+  import torch_npu
+  if torch.npu.is_available() and not torch.cuda.is_available():
+      enable_cuda_to_npu_shim()
+except: 
+  print("exception")
+
+# 1) 可选：如果你的权重来自 lightning 的 ckpt，放行其类（仅在可信来源时）
+try:
+    from torch.serialization import add_safe_globals
+    from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+    add_safe_globals([ModelCheckpoint])
+except Exception:
+    pass
+
+# 2) 统一把 torch.load 默认映射到 CPU，避免 CUDA 反序列化错误
+_orig_load = torch.load
+def _load_map_to_cpu(*args, **kwargs):
+    kwargs.setdefault("map_location", "cpu")
+    kwargs.setdefault("weights_only", False)
+    return _orig_load(*args, **kwargs)
+torch.load = _load_map_to_cpu
+
+

run_in_docker_ascend.sh

@@ -0,0 +1,4 @@
+#! /usr/bin/env bash
+image=harbor-contest.4pd.io/zhanghao/t2v:ascend-0.1
+device=0
+docker run -it -v `pwd`:/host -e ASCEND_VISIBLE_DEVICES=$device --device /dev/davinci$device:/dev/davinci0 --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc -v /mnt:/mnt -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info -v /etc/ascend_install.info:/etc/ascend_install.info --privileged --entrypoint bash  $image