use latest torch

pyproject.toml

@@ -11,15 +11,4 @@ dependencies = [
     "soundfile>=0.12.1",
     "torch>=2.5.1",
     "vocos>=0.1.0",
-]
-
-
-[[tool.uv.index]]
-name = "pytorch-cpu"
-url = "https://download.pytorch.org/whl/cpu"
-explicit = true
-
-[tool.uv.sources]
-torch = [
-  { index = "pytorch-cpu", marker = "platform_system != 'Darwin'"},
 ]

requirements.txt

@@ -1,9 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile pyproject.toml -o requirements.txt --python-platform x86_64-unknown-linux-gnu --emit-index-url --extra-index-url https://download.pytorch.org/whl/cu113
---index-url https://pypi.org/simple
---extra-index-url https://download.pytorch.org/whl/cu113
-
-
+#    uv pip compile pyproject.toml -o requirements.txt --python-platform x86_64-unknown-linux-gnu
 aiofiles==23.2.1
     # via gradio
 annotated-types==0.7.0
@@ -43,6 +39,7 @@ filelock==3.13.1
     # via
     #   huggingface-hub
     #   torch
+    #   triton
 fsspec==2024.2.0
     # via
     #   gradio-client
@@ -109,6 +106,38 @@ numpy==1.26.3
     #   pandas
     #   scipy
     #   vocos
+nvidia-cublas-cu12==12.4.5.8
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti-cu12==12.4.127
+    # via torch
+nvidia-cuda-nvrtc-cu12==12.4.127
+    # via torch
+nvidia-cuda-runtime-cu12==12.4.127
+    # via torch
+nvidia-cudnn-cu12==9.1.0.70
+    # via torch
+nvidia-cufft-cu12==11.2.1.3
+    # via torch
+nvidia-curand-cu12==10.3.5.147
+    # via torch
+nvidia-cusolver-cu12==11.6.1.9
+    # via torch
+nvidia-cusparse-cu12==12.3.1.170
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-nccl-cu12==2.21.5
+    # via torch
+nvidia-nvjitlink-cu12==12.4.127
+    # via
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+    #   torch
+nvidia-nvtx-cu12==12.4.127
+    # via torch
 orjson==3.10.12
     # via gradio
 packaging==24.1
@@ -175,13 +204,13 @@ sympy==1.13.1
     # via torch
 tomlkit==0.13.2
     # via gradio
-torch
+torch==2.5.1
     # via
     #   diffusion-speech-360h (pyproject.toml)
     #   encodec
     #   torchaudio
     #   vocos
-torchaudio
+torchaudio==2.5.1
     # via
     #   encodec
     #   vocos
@@ -189,6 +218,8 @@ tqdm==4.66.5
     # via
     #   huggingface-hub
     #   nltk
+triton==3.1.0
+    # via torch
 typeguard==4.2.0
     # via inflect
 typer==0.15.1

sample.py

@@ -174,11 +174,12 @@ def sample(
             embedding_vocab_size=model_config["embedding_vocab_size"],
             learn_sigma=model_config["learn_sigma"],
             in_channels=data_config["data_dim"],
-        ).to(device)
+        ).to(device).bfloat16
 
         state_dict = find_model(ckpt_path)
         model.load_state_dict(state_dict)
         model.eval()  # important!
+        model = model.bfloat16()
         model_cache[ckpt_path] = model
     else:
         model = model_cache[ckpt_path]
@@ -206,15 +207,17 @@ def sample(
         attn_mask=attn_mask,
     )
 
-    samples = diffusion.p_sample_loop(
-        model.forward_with_cfg,
-        z.shape,
-        z,
-        clip_denoised=False,
-        model_kwargs=model_kwargs,
-        progress=True,
-        device=device,
-    )
+    with torch.no_grad():
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            samples = diffusion.p_sample_loop(
+                model.forward_with_cfg,
+                z.shape,
+                z,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+                progress=True,
+                device=device,
+            )
     samples = [s.chunk(2, dim=0)[0] for s in samples]  # Remove null class samples
     return samples