diff --git a/comfy/utils.py b/comfy/utils.py index b35062824..c901347c4 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -50,7 +50,16 @@ def load_torch_file(ckpt, safe_load=False, device=None): if device is None: device = torch.device("cpu") if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"): - sd = safetensors.torch.load_file(ckpt, device=device.type) + try: + sd = safetensors.torch.load_file(ckpt, device=device.type) + except Exception as e: + if len(e.args) > 0: + message = e.args[0] + if "HeaderTooLarge" in message: + raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt or invalid. Make sure this is actually a safetensors file and not a ckpt or pt or other filetype.".format(message, ckpt)) + if "MetadataIncompleteBuffer" in message: + raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is incomplete. Check the file size and make sure you have copied/downloaded it correctly.".format(message, ckpt)) + raise e else: if safe_load or ALWAYS_SAFE_LOAD: pl_sd = torch.load(ckpt, map_location=device, weights_only=True) diff --git a/latent_preview.py b/latent_preview.py index 07f9cc68e..95d3cb733 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -12,7 +12,10 @@ MAX_PREVIEW_RESOLUTION = args.preview_size def preview_to_image(latent_image): latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1 .mul(0xFF) # to 0..255 - ).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device)) + ) + if comfy.model_management.directml_enabled: + latents_ubyte = latents_ubyte.to(dtype=torch.uint8) + latents_ubyte = latents_ubyte.to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device)) return Image.fromarray(latents_ubyte.numpy())