From 035a370fa1a6f02c1f2efdf4062a77d547da1b81 Mon Sep 17 00:00:00 2001 From: Chris Chance Date: Tue, 15 Oct 2024 14:53:54 -0400 Subject: [PATCH] disable autocast if MPS is available AMP Autocast not available on current pytorch, will be added in 2.5+ version and nightlies but those are buggy, for now just disable AMP if we're on MPS. --- nodes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index 59fce7d..2d0a702 100644 --- a/nodes.py +++ b/nodes.py @@ -1116,7 +1116,9 @@ class CogVideoXFunSampler: generator= torch.Generator(device="cpu").manual_seed(seed) - autocastcondition = not pipeline["onediff"] + # Note: if we're on MPS return False, otherwise check for onediff as normal + # TODO: Torch 2.5+ is supposed to add AMP support for MPS ... not available currently. + autocastcondition = not pipeline["onediff"] if not torch.backends.mps.is_available() else False autocast_context = torch.autocast(mm.get_autocast_device(device)) if autocastcondition else nullcontext() with autocast_context: video_length = int((video_length - 1) // pipe.vae.config.temporal_compression_ratio * pipe.vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1