Update inference.py

inference.py

@@ -18,7 +18,7 @@ SampleFn = Callable[
 ]
 ### Loading
 
-def load_model_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cuda', dtype: torch.dtype | None = torch.bfloat16, compile: bool = False, token: str | None = None) -> EchoDiT:
+def load_model_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cpu', dtype: torch.dtype | None = torch.bfloat16, compile: bool = False, token: str | None = None) -> EchoDiT:
     with torch.device('meta'):
         model = EchoDiT(
             latent_size=80, model_size=2048, num_layers=24, num_heads=16,
@@ -50,7 +50,7 @@ def load_model_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cu
 
     return model
     
-def load_fish_ae_from_hf(repo_id: str = 'jordand/fish-s1-dac-min', device: str = 'cuda', dtype: torch.dtype | None = torch.float32, compile: bool = False, token: str | None = None) -> DAC:
+def load_fish_ae_from_hf(repo_id: str = 'jordand/fish-s1-dac-min', device: str = 'cpu', dtype: torch.dtype | None = torch.float32, compile: bool = False, token: str | None = None) -> DAC:
     # have not tested lower precisions with fish AE yet
    
     with torch.device('meta'):
@@ -81,7 +81,7 @@ class PCAState:
     pca_mean: torch.Tensor
     latent_scale: float
 
-def load_pca_state_from_hf(repo_id: str = 'jordand/echo-tts', device: str = 'cuda', filename: str = 'pca_state.safetensors', token: str | None = None) -> 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)
     return PCAState(