From 19d08550f871b97534e586cbbce48c9075fde6b9 Mon Sep 17 00:00:00 2001
From: David <48716362+oldskool978@users.noreply.github.com>
Date: Thu, 13 Feb 2025 20:11:50 -0500
Subject: [PATCH] moved VAE so weights and tensors reside on the same hardware.

Moved VAE to device
fixes
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same.
---
 sonic_node.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/sonic_node.py b/sonic_node.py
index c9b644f..0fed903 100644
--- a/sonic_node.py
+++ b/sonic_node.py
@@ -240,6 +240,9 @@ class SONICSampler:
 
     def sampler_main(self, model, data_dict, seed, inference_steps, dynamic_scale, fps):
         print("***********Start infer  ***********")
+        vae = data_dict["vae"]
+        vae.first_stage_model = vae.first_stage_model.to(device)
+        vae.output_device = device
         iamge = model.process(data_dict["audio_tensor_list"],
                               data_dict["uncond_audio_tensor_list"],
                               data_dict["motion_buckets"],
@@ -248,7 +251,7 @@ class SONICSampler:
                               image_embeds=data_dict["image_embeddings"],
                               img_latent=data_dict["img_latent"],
                               fps=fps,
-                              vae=data_dict["vae"],
+                              vae=vae,
                               inference_steps=inference_steps,
                               dynamic_scale=dynamic_scale,
                               seed=seed