Upload app.py with huggingface_hub

app.py

@@ -4,61 +4,148 @@ os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
 
 import spaces
 import torch
-print(f"After import spaces: cuda_available={torch.cuda.is_available()}")
+import numpy as np
+import gradio as gr
+
+import subprocess
+import sys
 
-import subprocess, sys
 subprocess.run(
     [sys.executable, "-m", "pip", "install", "--no-deps",
      "chatterbox-tts==0.1.7", "chatterbox-flash==0.1.0"],
     check=True,
 )
-print(f"After pip install: cuda_available={torch.cuda.is_available()}")
 
-# Test: just import chatterbox, don't load model
 from chatterbox_flash import ChatterboxFlashTTS
-print(f"After chatterbox import: cuda_available={torch.cuda.is_available()}")
-
-# Test: load just the voice encoder (small model)
-from chatterbox.models.ve.model import VoiceEncoder
-from safetensors.torch import load_file
-import tempfile, os as _os
-from huggingface_hub import hf_hub_download
-
-print("Downloading ve.safetensors...")
-ve_path = hf_hub_download("ResembleAI/chatterbox-flash", "ve.safetensors")
-ve = VoiceEncoder()
-ve.load_state_dict(load_file(ve_path))
-ve.to("cpu").eval()
-print(f"After ve load: cuda_available={torch.cuda.is_available()}")
-
-# Test: load just the t3 model
-from chatterbox_flash.t3 import ChatterboxFlashT3
-print("Downloading t3_flash.safetensors...")
-t3_path = hf_hub_download("ResembleAI/chatterbox-flash", "t3_flash.safetensors")
-t3 = ChatterboxFlashT3(drf_block_size=16)
-t3_state = load_file(t3_path)
-if "model" in t3_state:
-    t3_state = t3_state["model"][0]
-t3.load_state_dict(t3_state)
-t3.to(device="cpu", dtype=torch.bfloat16).eval()
-print(f"After t3 load: cuda_available={torch.cuda.is_available()}")
-
-# Test: load s3gen
-from chatterbox.models.s3gen.s3gen import S3Gen
-print("Downloading s3gen.safetensors...")
-s3gen_path = hf_hub_download("ResembleAI/chatterbox-flash", "s3gen.safetensors")
-s3gen = S3Gen(meanflow=True)
-s3gen.load_state_dict(load_file(s3gen_path), strict=False)
-s3gen.to("cpu").eval()
-print(f"After s3gen load: cuda_available={torch.cuda.is_available()}")
-
-print("All models loaded on CPU. Testing GPU...")
 
-import gradio as gr
+MODEL_ID = "ResembleAI/chatterbox-flash"
+_tts = None
+
+
+def get_tts():
+    global _tts
+    if _tts is None:
+        print(f"Loading Chatterbox-Flash from {MODEL_ID}...")
+        _tts = ChatterboxFlashTTS.from_pretrained(
+            MODEL_ID, device="cuda", dtype=torch.bfloat16,
+        )
+        print("Model loaded successfully.")
+    return _tts
+
+
+@spaces.GPU(duration=120)
+def generate_tts(
+    text_input: str,
+    audio_prompt_path: str | None = None,
+    exaggeration: float = 0.5,
+    temperature: float = 0.6,
+    cfg_scale: float = 1.0,
+    num_steps: int = 10,
+    seed_num: int = 0,
+):
+    """Generate speech from text using Chatterbox-Flash block-diffusion TTS."""
+    tts = get_tts()
+
+    if seed_num != 0:
+        torch.manual_seed(int(seed_num))
+        torch.cuda.manual_seed(int(seed_num))
+        np.random.seed(int(seed_num))
+
+    generate_kwargs = {
+        "exaggeration": exaggeration,
+        "temperature": temperature,
+        "cfg_scale": cfg_scale,
+        "num_steps": num_steps,
+        "backend": "torch",
+    }
+
+    if audio_prompt_path:
+        generate_kwargs["audio_prompt_path"] = audio_prompt_path
+
+    wav = tts.generate(text_input[:300], **generate_kwargs)
+    return (tts.sr, wav.squeeze(0).cpu().numpy())
+
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Chatterbox-Flash TTS
+        Prior-calibrated block-diffusion zero-shot TTS by Resemble AI.
+        Provide a reference audio clip to clone a voice, or generate with the default voice.
+
+        [Paper](https://huggingface.co/papers/2605.30748) · [Model](https://huggingface.co/ResembleAI/chatterbox-flash) · [GitHub](https://github.com/resemble-ai/chatterbox-flash)
+        """
+    )
+
+    with gr.Row(elem_id="col-container"):
+        with gr.Column(scale=3):
+            text = gr.Textbox(
+                value="Sometimes it's better to just let things slide, you know?",
+                label="Text to synthesize (max 300 chars)",
+                max_lines=5,
+            )
+            ref_wav = gr.Audio(
+                sources=["upload", "microphone"],
+                type="filepath",
+                label="Reference Audio (for voice cloning)",
+            )
+
+            with gr.Accordion("Advanced settings", open=False):
+                exaggeration = gr.Slider(
+                    0.25, 2.0, step=0.05,
+                    label="Exaggeration (0.5=neutral, higher=more expressive)",
+                    value=0.5,
+                )
+                temperature = gr.Slider(
+                    0.05, 2.0, step=0.05,
+                    label="Temperature",
+                    value=0.6,
+                )
+                cfg_scale = gr.Slider(
+                    0.2, 1.0, step=0.05,
+                    label="CFG Scale",
+                    value=1.0,
+                )
+                num_steps = gr.Slider(
+                    1, 30, step=1,
+                    label="Denoising Steps",
+                    value=10,
+                )
+                seed_num = gr.Number(
+                    value=0, label="Seed (0=random)", precision=0,
+                )
+
+            run_btn = gr.Button("Generate", variant="primary")
+
+        with gr.Column(scale=2):
+            audio_output = gr.Audio(label="Output Audio")
+
+    gr.Examples(
+        examples=[
+            ["Sometimes it's better to just let things slide, you know?"],
+            ["The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs."],
+            ["In the depths of winter, I finally learned that within me lay an invincible summer."],
+        ],
+        inputs=[text],
+        outputs=[audio_output],
+        fn=generate_tts,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
 
-@spaces.GPU(duration=60)
-def test_gpu(text):
-    return f"CUDA: {torch.cuda.is_available()}, devices: {torch.cuda.device_count()}"
+    run_btn.click(
+        fn=generate_tts,
+        inputs=[
+            text, ref_wav, exaggeration,
+            temperature, cfg_scale, num_steps, seed_num,
+        ],
+        outputs=[audio_output],
+        api_name="generate",
+    )
 
-demo = gr.Interface(fn=test_gpu, inputs=gr.Textbox(value="test"), outputs=gr.Textbox())
-demo.launch(mcp_server=True)
+demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)