making pocket tts compat with cuda generation (#132)

Co-authored-by: gabrieldemarmiesse <gabrieldemarmiesse@gmail.com>

pocket_tts/modules/conv.py

@@ -85,8 +85,9 @@ class StreamingConv1d(StatefulModule):
         stride = self._stride
         # Effective kernel size accounting for dilation.
         kernel = self._effective_kernel_size
-        previous = torch.zeros(batch_size, self.conv.in_channels, kernel - stride)
-        first = torch.ones(batch_size, dtype=torch.bool)
+        device = self.conv.weight.device
+        previous = torch.zeros(batch_size, self.conv.in_channels, kernel - stride, device=device)
+        first = torch.ones(batch_size, dtype=torch.bool, device=device)
         return dict(previous=previous, first=first)
 
     def forward(self, x, model_state: dict | None):
@@ -144,7 +145,8 @@ class StreamingConvTranspose1d(StatefulModule):
     def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
         K = self._kernel_size
         S = self._stride
-        return dict(partial=torch.zeros(batch_size, self.convtr.out_channels, K - S))
+        device = self.convtr.weight.device
+        return dict(partial=torch.zeros(batch_size, self.convtr.out_channels, K - S, device=device))
 
     def forward(self, x, mimi_state: dict):
         layer_state = self.get_state(mimi_state)["partial"]

pocket_tts/modules/mimi_transformer.py

@@ -79,9 +79,12 @@ class MimiStreamingMultiheadAttention(StatefulModule):
         dim_per_head = self.embed_dim // self.num_heads
 
         state = {}
-        state["offset"] = torch.zeros(batch_size, dtype=torch.long)
-        state["cache"] = torch.zeros((2, batch_size, self.num_heads, sequence_length, dim_per_head))
-        state["end_offset"] = torch.zeros(batch_size, dtype=torch.long)
+        device = self.in_proj.weight.device
+        state["offset"] = torch.zeros(batch_size, dtype=torch.long, device=device)
+        state["cache"] = torch.zeros(
+            (2, batch_size, self.num_heads, sequence_length, dim_per_head), device=device
+        )
+        state["end_offset"] = torch.zeros(batch_size, dtype=torch.long, device=device)
         return state
 
     def increment_step(self, state, increment: int = 1):