update

OmniGen/model.py

@@ -197,7 +197,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
         model = cls(config)
         if os.path.exists(os.path.join(model_name, 'model.safetensors')):
             print("Loading safetensors")
-            ckpt = load_file(os.path.join(model_name, 'model.safetensors'), device='cpu')
+            ckpt = load_file(os.path.join(model_name, 'model.safetensors'))
         else:
             ckpt = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
         model.load_state_dict(ckpt)