diff --git a/comfy/extra_samplers/uni_pc.py b/comfy/extra_samplers/uni_pc.py index 5b80a8aff..c57e081e4 100644 --- a/comfy/extra_samplers/uni_pc.py +++ b/comfy/extra_samplers/uni_pc.py @@ -661,7 +661,7 @@ class UniPC: if x_t is None: if use_predictor: - pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s) + pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0])) # torch.einsum('k,bkchw->bchw', rhos_p, D1s) else: pred_res = 0 x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * pred_res @@ -669,7 +669,7 @@ class UniPC: if use_corrector: model_t = self.model_fn(x_t, t) if D1s is not None: - corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s) + corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0])) # torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s) else: corr_res = 0 D1_t = (model_t - model_prev_0) diff --git a/comfy/utils.py b/comfy/utils.py index bcefa1804..b35062824 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -43,7 +43,8 @@ if hasattr(torch.serialization, "add_safe_globals"): # TODO: this was added in torch.serialization.add_safe_globals([ModelCheckpoint, scalar, dtype, Float64DType, encode]) ALWAYS_SAFE_LOAD = True logging.info("Checkpoint files will always be loaded safely.") - +else: + logging.info("Warning, you are using an old pytorch version and some ckpt/pt files might be loaded unsafely. Upgrading to 2.4 or above is recommended.") def load_torch_file(ckpt, safe_load=False, device=None): if device is None: