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.
This commit is contained in:
Chris Chance 2024-10-15 14:53:54 -04:00 committed by GitHub
parent 09ed641575
commit 035a370fa1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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