Merge pull request #359 from KohakuBlueleaf/fix-dtype-error

Fix dtype mismatch
2024-03-09 15:38:00 +08:00 · 2024-03-09 15:38:00 +08:00 · 76cde3c2bf
parent 0086a19ca4 c97cfd1e83
commit 76cde3c2bf
1 changed files with 2 additions and 1 deletions
--- a/scripts/tilevae.py
+++ b/scripts/tilevae.py
@ -508,6 +508,7 @@ class VAEHook:
        @return: image
        """
        device = next(self.net.parameters()).device
+        dtype = next(self.net.parameters()).dtype
        net = self.net
        tile_size = self.tile_size
        is_decoder = self.is_decoder
@ -647,7 +648,7 @@ class VAEHook:

        # Done!
        pbar.close()
-        return result if result is not None else result_approx.to(device)
+        return result.to(dtype) if result is not None else result_approx.to(device, dtype=dtype)


 class Script(scripts.Script):