fix seed, clean code

app.py

@@ -1,3 +1,15 @@
+from PIL import Image, ImageDraw
+from functools import lru_cache
+import gradio as gr
+from smooth_cache_helper import SmoothCacheHelper
+from f5_tts.infer.utils_infer import (
+    cross_fade_duration,
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    speed
+)
 import os
 from importlib.resources import files
 
@@ -13,20 +25,6 @@ import sys
 sys.path.append('F5-TTS/src')
 sys.path.append('SmoothCache/SmoothCache')
 
-from f5_tts.infer.utils_infer import (
-    cross_fade_duration,
-    infer_process,
-    load_model,
-    load_vocoder,
-    preprocess_ref_audio_text,
-    speed
-)
-from smooth_cache_helper import SmoothCacheHelper
-
-import gradio as gr
-import numpy as np
-from functools import lru_cache
-from PIL import Image, ImageDraw
 
 try:
     import spaces
@@ -35,12 +33,14 @@ try:
 except ImportError:
     USING_SPACES = False
 
+
 def gpu_decorator(func):
     if USING_SPACES:
         return spaces.GPU(func)
     else:
         return func
 
+
 # Constants
 layer_names = ['attn', 'ff']
 colors_rgb = [(255, 103, 35), (0, 210, 106)]  # orange, green
@@ -76,14 +76,19 @@ cache_schedule = {
     'ff':   presets[default_preset]['ff'][:]
 }
 
-config = tomli.load(open(os.path.join(files("f5_tts").joinpath("infer/examples/basic"), "basic.toml"), "rb"))
+seed = np.random.randint(0, 2**31 - 1)
+torch.manual_seed(seed)
+
+config = tomli.load(open(os.path.join(files("f5_tts").joinpath(
+    "infer/examples/basic"), "basic.toml"), "rb"))
 
 model = config.get("model", "F5TTS_v1_Base")
 ckpt_file = config.get("ckpt_file", "")
 vocab_file = config.get("vocab_file", "")
 
 model_cfg = OmegaConf.load(
-    config.get("model_cfg", str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+    config.get("model_cfg", str(
+        files("f5_tts").joinpath(f"configs/{model}.yaml")))
 )
 model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
 model_arc = model_cfg.model.arch
@@ -91,7 +96,8 @@ model_arc = model_cfg.model.arch
 repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
 
 if not ckpt_file:
-    ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}"))
+    ckpt_file = str(cached_path(
+        f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}"))
 
 if not vocab_file:
     vocab_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/vocab.txt"))
@@ -103,10 +109,12 @@ ema_model = load_model(
 
 vocoder = load_vocoder()
 
+
 @gpu_decorator
 def render_grid(schedule: dict) -> np.ndarray:
     n_steps = len(schedule['attn'])
-    img = Image.new("RGB", (n_steps * (cell_size + spacing), n_layers * (cell_size + spacing)), "white")
+    img = Image.new("RGB", (n_steps * (cell_size + spacing),
+                    n_layers * (cell_size + spacing)), "white")
     draw = ImageDraw.Draw(img)
 
     for row in range(n_layers):
@@ -121,6 +129,7 @@ def render_grid(schedule: dict) -> np.ndarray:
 
     return np.array(img)
 
+
 @gpu_decorator
 def apply_preset(preset_name):
     global cache_schedule
@@ -130,6 +139,7 @@ def apply_preset(preset_name):
         cache_schedule['ff'] = schedule['ff'][:]
     return render_grid(cache_schedule), len(cache_schedule['attn'])
 
+
 @gpu_decorator
 def toggle_cell(evt: gr.SelectData):
     global cache_schedule
@@ -140,6 +150,7 @@ def toggle_cell(evt: gr.SelectData):
         cache_schedule[layer][col] ^= 1
     return render_grid(cache_schedule), "Custom"
 
+
 @gpu_decorator
 def reset_schedule(n_steps):
     global cache_schedule
@@ -149,46 +160,37 @@ def reset_schedule(n_steps):
     }
     return render_grid(cache_schedule), "Custom"
 
+
 @gpu_decorator
 def update_nfe(nfe_value):
     return reset_schedule(nfe_value)
 
+
 @gpu_decorator
 def load_default():
     return render_grid(cache_schedule), default_preset
 
+
 @lru_cache(maxsize=1000)  # NOTE. need to ensure params of infer() hashable
 @gpu_decorator
 def infer(
     ref_audio_orig,
     ref_text,
     gen_text,
-    #model,
-    #remove_silence,
-    #seed,
-    #cross_fade_duration=0.15,
     nfe_step=32,
-    #speed=1,
-    #show_info=gr.Info,
 ):
     global cache_schedule
-    show_info=gr.Info
+    show_info = gr.Info
     if not ref_audio_orig:
         gr.Warning("Please provide reference audio.")
         return gr.update(), gr.update(), ref_text
 
-    # Set inference seed
-    # if seed < 0 or seed > 2**31 - 1:
-    #     gr.Warning("Seed must in range 0 ~ 2147483647. Using random seed instead.")
-    seed = np.random.randint(0, 2**31 - 1)
-    torch.manual_seed(seed)
-    used_seed = seed
-
     if not gen_text.strip():
         gr.Warning("Please enter text to generate or upload a text file.")
         return gr.update(), gr.update(), ref_text
 
-    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
+    ref_audio, ref_text = preprocess_ref_audio_text(
+        ref_audio_orig, ref_text, show_info=show_info)
     start_time = time.time()
     final_wave, final_sample_rate, _ = infer_process(
         ref_audio,
@@ -206,7 +208,7 @@ def infer(
     cache_helper = SmoothCacheHelper(
         model=ema_model.transformer,
         block_classes=get_class("f5_tts.model.modules.DiTBlock"),
-        components_to_wrap=['attn','ff'],
+        components_to_wrap=['attn', 'ff'],
         schedule=cache_schedule
     )
     cache_helper.enable()
@@ -226,79 +228,46 @@ def infer(
     process_time_cache = time.time() - start_time
     cache_helper.disable()
 
-    # Remove silence
-    # if remove_silence:
-    #     with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
-    #         temp_path = f.name
-    #     try:
-    #         sf.write(temp_path, final_wave, final_sample_rate)
-    #         remove_silence_for_generated_wav(f.name)
-    #         final_wave, _ = torchaudio.load(f.name)
-    #     finally:
-    #         os.unlink(temp_path)
-    #     final_wave = final_wave.squeeze().cpu().numpy()
-
-    # Save the spectrogram
-    # with tempfile.NamedTemporaryFile(suffix=".png", **tempfile_kwargs) as tmp_spectrogram:
-    #     spectrogram_path = tmp_spectrogram.name
-    # save_spectrogram(combined_spectrogram, spectrogram_path)
-
     return (final_sample_rate, final_wave), (final_sample_rate_cache, final_wave_cache), process_time, process_time_cache
 
 
 with gr.Blocks() as demo:
     gr.Markdown("## F5-TTS + SmoothCache")
     ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
-    ref_text_input = gr.Textbox(
-        label="Reference Text",
-        #info="Leave blank to automatically transcribe the reference audio. If you enter text or upload a file, it will override automatic transcription.",
-        # lines=2,
-        # scale=4,
-    )
-    gen_text_input = gr.Textbox(
-        label="Text to Generate",
-        # lines=10,
-        # max_lines=40,
-        # scale=4,
-    )
+    ref_text_input = gr.Textbox(label="Reference Text")
+    gen_text_input = gr.Textbox(label="Text to Generate")
     with gr.Row():
         with gr.Column(scale=0):
-            preset_dropdown = gr.Dropdown(choices=list(presets.keys()) + ["Custom"], label="Choose Preset", value=default_preset)
-            nfe_slider = gr.Slider(4, 64, value=len(cache_schedule['attn']), step=1, label="Number of Steps (NFE)")
+            preset_dropdown = gr.Dropdown(choices=list(
+                presets.keys()) + ["Custom"], label="Choose Preset", value=default_preset)
+            nfe_slider = gr.Slider(4, 64, value=len(
+                cache_schedule['attn']), step=1, label="Number of Steps (NFE)")
         with gr.Column(scale=1):
-            gr.Markdown("Click Grid to Customize Cache Schedule<br>🟧 = Compute Attn Layer / 🟩 = Compute FFN Layer / ⬜ = Cached Layer")
+            gr.Markdown(
+                "Click Grid to Customize Cache Schedule<br>🟧 = Compute Attn Layer / 🟩 = Compute FFN Layer / ⬜ = Cached Layer")
             image = gr.Image(type="numpy", label="", interactive=True, scale=1)
-        #reset_btn = gr.Button("Reset to All Cached")
-    #current_label = gr.Textbox(label="Current Preset", interactive=False)
     generate_btn = gr.Button("Synthesize", variant="primary")
     with gr.Row():
         with gr.Group():
             audio_output = gr.Audio(label="Synthesized Audio (No Cache)")
-            process_time = gr.Textbox(label="⏱ Process Time", interactive=False)
+            process_time = gr.Textbox(
+                label="⏱ Process Time", interactive=False)
         with gr.Group():
             audio_output_cache = gr.Audio(label="Synthesized Audio (Cache)")
-            process_time_cache = gr.Textbox(label="⏱ Process Time", interactive=False)
+            process_time_cache = gr.Textbox(
+                label="⏱ Process Time", interactive=False)
 
     # Wire up logic
-    preset_dropdown.change(fn=apply_preset, inputs=preset_dropdown, outputs=[image, nfe_slider])
-    #preset_dropdown.change(fn=lambda x: x, inputs=preset_dropdown, outputs=current_label)
+    preset_dropdown.change(
+        fn=apply_preset, inputs=preset_dropdown, outputs=[image, nfe_slider])
     image.select(fn=toggle_cell, outputs=[image, preset_dropdown])
-    #reset_btn.click(fn=reset_schedule, inputs=nfe_slider, outputs=[image, preset_dropdown])
-    nfe_slider.input(fn=update_nfe, inputs=nfe_slider, outputs=[image, preset_dropdown])
+    nfe_slider.input(fn=update_nfe, inputs=nfe_slider,
+                     outputs=[image, preset_dropdown])
     generate_btn.click(
         infer,
-        inputs=[
-            ref_audio_input,
-            ref_text_input,
-            gen_text_input,
-            #remove_silence,
-            #randomize_seed,
-            #np.random.randint(0, 2**31 - 1),
-            #cross_fade_duration_slider,
-            nfe_slider,
-            #speed_slider,
-        ],
-        outputs=[audio_output, audio_output_cache, process_time, process_time_cache],
+        inputs=[ref_audio_input, ref_text_input, gen_text_input, nfe_slider],
+        outputs=[audio_output, audio_output_cache,
+                 process_time, process_time_cache],
     )
     demo.load(fn=load_default, outputs=[image, preset_dropdown])