v1.1.9

* fix raw.arrow missing rows

---------

Co-authored-by: SWivid <swivid@qq.com>

.github/workflows/sync-hf.yaml

@@ -1,17 +0,0 @@
-name: Sync to HF Space
-
-on:
-  release:
-    types: [published]
-
-jobs:
-  trigger_curl:
-    runs-on: ubuntu-latest
-
-    steps:
-    - name: Send cURL POST request
-      run: |
-        curl -X POST https://mrfakename-sync-f5.hf.space/gradio_api/call/refresh \
-        -s \
-        -H "Content-Type: application/json" \
-        -d "{\"data\": [\"${{ secrets.REFRESH_PASSWORD }}\"]}"

README.md

@@ -2,11 +2,12 @@
 
 [![python](https://img.shields.io/badge/Python-3.10-brightgreen)](https://github.com/SWivid/F5-TTS)
 [![arXiv](https://img.shields.io/badge/arXiv-2410.06885-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2410.06885)
-[![demo](https://img.shields.io/badge/GitHub-Demo%20page-orange.svg)](https://swivid.github.io/F5-TTS/)
-[![hfspace](https://img.shields.io/badge/🤗-Space%20demo-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
-[![msspace](https://img.shields.io/badge/🤖-Space%20demo-blue)](https://modelscope.cn/studios/modelscope/E2-F5-TTS)
-[![lab](https://img.shields.io/badge/X--LANCE-Lab-grey?labelColor=lightgrey)](https://x-lance.sjtu.edu.cn/)
-[![lab](https://img.shields.io/badge/Peng%20Cheng-Lab-grey?labelColor=lightgrey)](https://www.pcl.ac.cn)
+[![demo](https://img.shields.io/badge/GitHub-Demo-orange.svg)](https://swivid.github.io/F5-TTS/)
+[![hfspace](https://img.shields.io/badge/🤗-HF%20Space-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+[![msspace](https://img.shields.io/badge/🤖-MS%20Space-blue)](https://modelscope.cn/studios/AI-ModelScope/E2-F5-TTS)
+[![lab](https://img.shields.io/badge/🏫-X--LANCE-grey?labelColor=lightgrey)](https://x-lance.sjtu.edu.cn/)
+[![lab](https://img.shields.io/badge/🏫-SII-grey?labelColor=lightgrey)](https://www.sii.edu.cn/)
+[![lab](https://img.shields.io/badge/🏫-PCL-grey?labelColor=lightgrey)](https://www.pcl.ac.cn)
 <!-- <img src="https://github.com/user-attachments/assets/12d7749c-071a-427c-81bf-b87b91def670" alt="Watermark" style="width: 40px; height: auto"> -->
 
 **F5-TTS**: Diffusion Transformer with ConvNeXt V2, faster trained and inference.
@@ -26,8 +27,8 @@
 ### Create a separate environment if needed
 
 ```bash
-# Create a python 3.10 conda env (you could also use virtualenv)
-conda create -n f5-tts python=3.10
+# Create a conda env with python_version>=3.10  (you could also use virtualenv)
+conda create -n f5-tts python=3.11
 conda activate f5-tts
 ```
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "f5-tts"
-version = "1.1.6"
+version = "1.1.9"
 description = "F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching"
 readme = "README.md"
 license = {text = "MIT License"}
@@ -14,18 +14,18 @@ classifiers = [
     "Programming Language :: Python :: 3",
 ]
 dependencies = [
-    "accelerate>=0.33.0,!=1.7.0",
-    "bitsandbytes>0.37.0; platform_machine != 'arm64' and platform_system != 'Darwin'",
+    "accelerate>=0.33.0",
+    "bitsandbytes>0.37.0; platform_machine!='arm64' and platform_system!='Darwin'",
     "cached_path",
     "click",
     "datasets",
     "ema_pytorch>=0.5.2",
-    "gradio>=3.45.2",
+    "gradio>=5.0.0",
     "hydra-core>=1.3.0",
     "jieba",
     "librosa",
     "matplotlib",
-    "numpy<=1.26.4",
+    "numpy<=1.26.4; python_version<='3.10'",
     "pydantic<=2.10.6",
     "pydub",
     "pypinyin",

src/f5_tts/infer/infer_gradio.py

@@ -943,9 +943,9 @@ with gr.Blocks() as app_credits:
 with gr.Blocks() as app:
     gr.Markdown(
         f"""
-# E2/F5 TTS
+# F5-TTS Demo Space
 
-This is {"a local web UI for [F5 TTS](https://github.com/SWivid/F5-TTS)" if not USING_SPACES else "an online demo for [F5-TTS](https://github.com/SWivid/F5-TTS)"} with advanced batch processing support. This app supports the following TTS models:
+This is {"a local web UI for [F5-TTS](https://github.com/SWivid/F5-TTS)" if not USING_SPACES else "an online demo for [F5-TTS](https://github.com/SWivid/F5-TTS)"} with advanced batch processing support. This app supports the following TTS models:
 
 * [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
 * [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)

src/f5_tts/model/backbones/dit.py

@@ -29,11 +29,16 @@ from f5_tts.model.modules import (
 
 
 class TextEmbedding(nn.Module):
-    def __init__(self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2):
+    def __init__(
+        self, text_num_embeds, text_dim, mask_padding=True, average_upsampling=False, conv_layers=0, conv_mult=2
+    ):
         super().__init__()
         self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
 
         self.mask_padding = mask_padding  # mask filler and batch padding tokens or not
+        self.average_upsampling = average_upsampling  # zipvoice-style text late average upsampling (after text encoder)
+        if average_upsampling:
+            assert mask_padding, "text_embedding_average_upsampling requires text_mask_padding to be True"
 
         if conv_layers > 0:
             self.extra_modeling = True
@@ -45,11 +50,47 @@ class TextEmbedding(nn.Module):
         else:
             self.extra_modeling = False
 
-    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
+    def average_upsample_text_by_mask(self, text, text_mask, audio_mask):
+        batch, text_len, text_dim = text.shape
+
+        if audio_mask is None:
+            audio_mask = torch.ones_like(text_mask, dtype=torch.bool)
+        valid_mask = audio_mask & text_mask
+        audio_lens = audio_mask.sum(dim=1)  # [batch]
+        valid_lens = valid_mask.sum(dim=1)  # [batch]
+
+        upsampled_text = torch.zeros_like(text)
+
+        for i in range(batch):
+            audio_len = audio_lens[i].item()
+            valid_len = valid_lens[i].item()
+
+            if valid_len == 0:
+                continue
+
+            valid_ind = torch.where(valid_mask[i])[0]
+            valid_data = text[i, valid_ind, :]  # [valid_len, text_dim]
+
+            base_repeat = audio_len // valid_len
+            remainder = audio_len % valid_len
+
+            indices = []
+            for j in range(valid_len):
+                repeat_count = base_repeat + (1 if j >= valid_len - remainder else 0)
+                indices.extend([j] * repeat_count)
+
+            indices = torch.tensor(indices[:audio_len], device=text.device, dtype=torch.long)
+            upsampled = valid_data[indices]  # [audio_len, text_dim]
+
+            upsampled_text[i, :audio_len, :] = upsampled
+
+        return upsampled_text
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False, audio_mask: bool["b n"] | None = None):  # noqa: F722
         text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
         text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
         batch, text_len = text.shape[0], text.shape[1]
-        text = F.pad(text, (0, seq_len - text_len), value=0)
+        text = F.pad(text, (0, seq_len - text_len), value=0)  # (opt.) if not self.average_upsampling:
         if self.mask_padding:
             text_mask = text == 0
 
@@ -61,7 +102,7 @@ class TextEmbedding(nn.Module):
         # possible extra modeling
         if self.extra_modeling:
             # sinus pos emb
-            batch_start = torch.zeros((batch,), dtype=torch.long)
+            batch_start = torch.zeros((batch,), device=text.device, dtype=torch.long)
             pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
             text_pos_embed = self.freqs_cis[pos_idx]
             text = text + text_pos_embed
@@ -75,6 +116,9 @@ class TextEmbedding(nn.Module):
             else:
                 text = self.text_blocks(text)
 
+        if self.average_upsampling:
+            text = self.average_upsample_text_by_mask(text, ~text_mask, audio_mask)
+
         return text
 
 
@@ -113,6 +157,7 @@ class DiT(nn.Module):
         text_num_embeds=256,
         text_dim=None,
         text_mask_padding=True,
+        text_embedding_average_upsampling=False,
         qk_norm=None,
         conv_layers=0,
         pe_attn_head=None,
@@ -127,7 +172,11 @@ class DiT(nn.Module):
         if text_dim is None:
             text_dim = mel_dim
         self.text_embed = TextEmbedding(
-            text_num_embeds, text_dim, mask_padding=text_mask_padding, conv_layers=conv_layers
+            text_num_embeds,
+            text_dim,
+            mask_padding=text_mask_padding,
+            average_upsampling=text_embedding_average_upsampling,
+            conv_layers=conv_layers,
         )
         self.text_cond, self.text_uncond = None, None  # text cache
         self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
@@ -190,19 +239,20 @@ class DiT(nn.Module):
         drop_audio_cond: bool = False,
         drop_text: bool = False,
         cache: bool = True,
+        audio_mask: bool["b n"] | None = None,  # noqa: F722
     ):
         seq_len = x.shape[1]
         if cache:
             if drop_text:
                 if self.text_uncond is None:
-                    self.text_uncond = self.text_embed(text, seq_len, drop_text=True)
+                    self.text_uncond = self.text_embed(text, seq_len, drop_text=True, audio_mask=audio_mask)
                 text_embed = self.text_uncond
             else:
                 if self.text_cond is None:
-                    self.text_cond = self.text_embed(text, seq_len, drop_text=False)
+                    self.text_cond = self.text_embed(text, seq_len, drop_text=False, audio_mask=audio_mask)
                 text_embed = self.text_cond
         else:
-            text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+            text_embed = self.text_embed(text, seq_len, drop_text=drop_text, audio_mask=audio_mask)
 
         x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
 
@@ -230,13 +280,19 @@ class DiT(nn.Module):
         # t: conditioning time, text: text, x: noised audio + cond audio + text
         t = self.time_embed(time)
         if cfg_infer:  # pack cond & uncond forward: b n d -> 2b n d
-            x_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache)
-            x_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache)
+            x_cond = self.get_input_embed(
+                x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache, audio_mask=mask
+            )
+            x_uncond = self.get_input_embed(
+                x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache, audio_mask=mask
+            )
             x = torch.cat((x_cond, x_uncond), dim=0)
             t = torch.cat((t, t), dim=0)
             mask = torch.cat((mask, mask), dim=0) if mask is not None else None
         else:
-            x = self.get_input_embed(x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache)
+            x = self.get_input_embed(
+                x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache, audio_mask=mask
+            )
 
         rope = self.rotary_embed.forward_from_seq_len(seq_len)
 

src/f5_tts/model/trainer.py

@@ -149,7 +149,7 @@ class Trainer:
         if self.is_main:
             checkpoint = dict(
                 model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
-                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
+                optimizer_state_dict=self.optimizer.state_dict(),
                 ema_model_state_dict=self.ema_model.state_dict(),
                 scheduler_state_dict=self.scheduler.state_dict(),
                 update=update,
@@ -242,7 +242,7 @@ class Trainer:
                     del checkpoint["model_state_dict"][key]
 
             self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
-            self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint["optimizer_state_dict"])
+            self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
             if self.scheduler:
                 self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
             update = checkpoint["update"]

src/f5_tts/train/datasets/prepare_csv_wavs.py

@@ -208,11 +208,11 @@ def save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_fine
     out_dir.mkdir(exist_ok=True, parents=True)
     print(f"\nSaving to {out_dir} ...")
 
-    # Save dataset with improved batch size for better I/O performance
     raw_arrow_path = out_dir / "raw.arrow"
-    with ArrowWriter(path=raw_arrow_path.as_posix(), writer_batch_size=100) as writer:
+    with ArrowWriter(path=raw_arrow_path.as_posix()) as writer:
         for line in tqdm(result, desc="Writing to raw.arrow ..."):
             writer.write(line)
+        writer.finalize()
 
     # Save durations to JSON
     dur_json_path = out_dir / "duration.json"

src/f5_tts/train/datasets/prepare_emilia.py

@@ -181,6 +181,7 @@ def main():
     with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer:
         for line in tqdm(result, desc="Writing to raw.arrow ..."):
             writer.write(line)
+        writer.finalize()
 
     # dup a json separately saving duration in case for DynamicBatchSampler ease
     with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f:

src/f5_tts/train/datasets/prepare_emilia_v2.py

@@ -68,6 +68,7 @@ def main():
     with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer:
         for line in tqdm(result, desc="Writing to raw.arrow ..."):
             writer.write(line)
+        writer.finalize()
 
     with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f:
         json.dump({"duration": duration_list}, f, ensure_ascii=False)

src/f5_tts/train/datasets/prepare_libritts.py

@@ -62,6 +62,7 @@ def main():
     with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer:
         for line in tqdm(result, desc="Writing to raw.arrow ..."):
             writer.write(line)
+        writer.finalize()
 
     # dup a json separately saving duration in case for DynamicBatchSampler ease
     with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f:

src/f5_tts/train/datasets/prepare_ljspeech.py

@@ -39,6 +39,7 @@ def main():
     with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer:
         for line in tqdm(result, desc="Writing to raw.arrow ..."):
             writer.write(line)
+        writer.finalize()
 
     # dup a json separately saving duration in case for DynamicBatchSampler ease
     with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f:

src/f5_tts/train/finetune_gradio.py

@@ -178,45 +178,12 @@ def get_audio_duration(audio_path):
     return audio.shape[1] / sample_rate
 
 
-def get_rms(
-    y,
-    frame_length=2048,
-    hop_length=512,
-    pad_mode="constant",
-):  # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py
-    padding = (int(frame_length // 2), int(frame_length // 2))
-    y = np.pad(y, padding, mode=pad_mode)
-
-    axis = -1
-    # put our new within-frame axis at the end for now
-    out_strides = y.strides + tuple([y.strides[axis]])
-    # Reduce the shape on the framing axis
-    x_shape_trimmed = list(y.shape)
-    x_shape_trimmed[axis] -= frame_length - 1
-    out_shape = tuple(x_shape_trimmed) + tuple([frame_length])
-    xw = np.lib.stride_tricks.as_strided(y, shape=out_shape, strides=out_strides)
-    if axis < 0:
-        target_axis = axis - 1
-    else:
-        target_axis = axis + 1
-    xw = np.moveaxis(xw, -1, target_axis)
-    # Downsample along the target axis
-    slices = [slice(None)] * xw.ndim
-    slices[axis] = slice(0, None, hop_length)
-    x = xw[tuple(slices)]
-
-    # Calculate power
-    power = np.mean(np.abs(x) ** 2, axis=-2, keepdims=True)
-
-    return np.sqrt(power)
-
-
 class Slicer:  # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py
     def __init__(
         self,
         sr: int,
         threshold: float = -40.0,
-        min_length: int = 2000,
+        min_length: int = 20000,  # 20 seconds
         min_interval: int = 300,
         hop_size: int = 20,
         max_sil_kept: int = 2000,
@@ -247,7 +214,7 @@ class Slicer:  # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.
             samples = waveform
         if samples.shape[0] <= self.min_length:
             return [waveform]
-        rms_list = get_rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
+        rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
         sil_tags = []
         silence_start = None
         clip_start = 0
@@ -301,8 +268,7 @@ class Slicer:  # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.
             silence_end = min(total_frames, silence_start + self.max_sil_kept)
             pos = rms_list[silence_start : silence_end + 1].argmin() + silence_start
             sil_tags.append((pos, total_frames + 1))
-        # Apply and return slices.
-        ####音频+起始时间+终止时间
+        # Apply and return slices: [chunk, start, end]
         if len(sil_tags) == 0:
             return [[waveform, 0, int(total_frames * self.hop_size)]]
         else:
@@ -830,9 +796,10 @@ def create_metadata(name_project, ch_tokenizer, progress=gr.Progress()):
     min_second = round(min(duration_list), 2)
     max_second = round(max(duration_list), 2)
 
-    with ArrowWriter(path=file_raw, writer_batch_size=1) as writer:
+    with ArrowWriter(path=file_raw) as writer:
         for line in progress.tqdm(result, total=len(result), desc="prepare data"):
             writer.write(line)
+        writer.finalize()
 
     with open(file_duration, "w") as f:
         json.dump({"duration": duration_list}, f, ensure_ascii=False)