Update inference.py

inference.py

@@ -57,14 +57,16 @@ def load_fish_ae_from_hf(repo_id: str = 'jordand/fish-s1-dac-min', device: str =
         fish_ae = build_ae()
 
     w_path = hf_hub_download(repo_id, 'pytorch_model.safetensors', token=token)
+    
+    # Загружаем на CPU, затем перемещаем на нужное устройство
+    state = st.load_file(w_path, device='cpu')
+    
     if dtype is not None and dtype != torch.float32:
-        state = st.load_file(w_path, device='cpu')
         state = {k: v.to(dtype=dtype) for k, v in state.items()}
-        state = {k: v.to(device=device) for k, v in state.items()}
-        fish_ae.load_state_dict(state, strict=False, assign=True)
-    else:
-        state = st.load_file(w_path, device=device)
-        fish_ae.load_state_dict(state, strict=False, assign=True)
+    
+    # Перемещаем на указанное устройство
+    state = {k: v.to(device=device) for k, v in state.items()}
+    fish_ae.load_state_dict(state, strict=False, assign=True)
 
     fish_ae = fish_ae.eval().to(device)  # Явно перемещаем модель на устройство
 
@@ -83,7 +85,13 @@ class PCAState:
 
 def load_pca_state_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cpu', filename: str = 'pca_state.safetensors', token: str | None = None) -> PCAState:
     p_path = hf_hub_download(repo_id, filename, token=token)
-    t = st.load_file(p_path, device=device)  # Загружаем напрямую на нужное устройство
+    
+    # Загружаем на CPU, затем перемещаем на нужное устройство
+    t = st.load_file(p_path, device='cpu')
+    
+    # Перемещаем тензоры на указанное устройство
+    t = {k: v.to(device=device) for k, v in t.items()}
+    
     return PCAState(
         pca_components=t["pca_components"],
         pca_mean=t["pca_mean"],