update

nested_attention_pipeline.py

@@ -4,6 +4,7 @@ from typing import List
 import torch
 from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+from safetensors import load_file
 
 from nested_attention_processor import AttnProcessor, NestedAttnProcessor
 from utils import get_generator
@@ -110,7 +111,7 @@ class NestedAdapterInference:
 
     def load_nested_adapter(self):
         state_dict = {"adapter_modules": {}, "qformer": {}}
-        f = torch.load(self.adapter_ckpt, map_location="cpu")
+        f = load_file(self.adapter_ckpt)
         for key in f.keys():
             if key.startswith("adapter_modules."):
                 state_dict["adapter_modules"][key.replace("adapter_modules.", "")] = f[