make compatible with comfy save/load latents

2024-10-24 03:23:01 +03:00 · 2024-10-24 03:23:01 +03:00 · f714748ad4
commit f714748ad4
parent 9bb3a79275
1 changed files with 2 additions and 0 deletions
--- a/nodes.py
+++ b/nodes.py
@ -266,6 +266,8 @@ class MochiDecode:
        device = mm.get_torch_device()
        offload_device = mm.unet_offload_device()
        samples = samples["samples"]
+        samples = samples.to(torch.bfloat16).to(device)
+
        
        def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
            blend_extent = min(a.shape[3], b.shape[3], blend_extent)