File size: 17,311 Bytes
240a407
5cbfdab
deb04b6
240a407
91d6893
240a407
 
189eecb
240a407
 
 
 
 
 
 
 
 
 
 
 
91d6893
 
 
deb04b6
74dbc75
91d6893
74dbc75
deb04b6
ee39abc
 
 
 
74dbc75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91d6893
74dbc75
 
 
 
 
 
91d6893
 
 
 
 
b32e924
74dbc75
91d6893
 
 
 
 
b32e924
91d6893
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74dbc75
 
 
91d6893
 
 
74dbc75
91d6893
 
74dbc75
91d6893
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240a407
91d6893
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240a407
74dbc75
 
ee39abc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240a407
74dbc75
240a407
 
 
74dbc75
240a407
74dbc75
240a407
189eecb
240a407
74dbc75
 
91d6893
 
240a407
 
74dbc75
240a407
 
74dbc75
240a407
 
 
ee39abc
 
 
 
 
 
 
 
240a407
ee39abc
74dbc75
 
 
 
 
 
 
 
91d6893
ee39abc
 
 
7be21d2
91d6893
 
 
 
 
 
240a407
 
 
 
91d6893
74dbc75
91d6893
 
 
 
 
ee39abc
91d6893
ee39abc
91d6893
 
 
 
 
 
ee39abc
91d6893
 
ee39abc
91d6893
 
 
 
240a407
deb04b6
240a407
deb04b6
91d6893
 
240a407
91d6893
240a407
 
 
 
 
 
91d6893
 
 
 
 
 
 
 
 
ee39abc
 
 
 
 
 
 
91d6893
240a407
91d6893
deb04b6
91d6893
 
74dbc75
 
 
 
91d6893
 
 
 
ee39abc
91d6893
ee39abc
91d6893
 
240a407
 
ee39abc
240a407
 
 
 
deb04b6
240a407
 
91d6893
5961c78
7be21d2
 
91d6893
7be21d2
 
 
5961c78
91d6893
7be21d2
91d6893
7be21d2
 
 
240a407
5961c78
 
74dbc75
240a407
91d6893
240a407
91d6893
240a407
91d6893
 
 
 
5961c78
240a407
5961c78
 
 
5cbfdab
deb04b6
 
74dbc75
 
 
 
ee39abc
 
 
 
 
 
 
 
74dbc75
 
 
 
ee39abc
74dbc75
ee39abc
 
74dbc75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ee39abc
74dbc75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5cbfdab
 
91d6893
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
import gradio as gr
import os
import traceback
import torch
import gc
from huggingface_hub import hf_hub_download
import shutil
import spaces

try:
    from config import MODEL_REPO_ID, MODEL_FILES, LOCAL_MODEL_PATH
except ImportError:
    MODEL_REPO_ID = "ramimu/chatterbox-voice-cloning-model"
    LOCAL_MODEL_PATH = "./chatterbox_model_files"
    MODEL_FILES = ["s3gen.pt", "t3_cfg.pt", "ve.pt", "tokenizer.json"]

try:
    from chatterbox.tts import ChatterboxTTS
    chatterbox_available = True
    print("Chatterbox TTS imported successfully")
except ImportError as e:
    print(f"Failed to import ChatterboxTTS: {e}")
    chatterbox_available = False

# Global model variable - will be loaded inside GPU function
model = None
model_loaded = False

# Text length limits for the model
MAX_CHARS_PER_GENERATION = 1000  # Safe limit for single generation
MAX_CHARS_TOTAL = 5000           # Maximum we'll accept via API

def download_model_files():
    """Download model files with error handling."""
    print(f"Checking for model files in {LOCAL_MODEL_PATH}...")
    os.makedirs(LOCAL_MODEL_PATH, exist_ok=True)
    
    for filename in MODEL_FILES:
        local_path = os.path.join(LOCAL_MODEL_PATH, filename)
        if not os.path.exists(local_path):
            print(f"Downloading {filename} from {MODEL_REPO_ID}...")
            try:
                downloaded_path = hf_hub_download(
                    repo_id=MODEL_REPO_ID,
                    filename=filename,
                    cache_dir="./cache",
                    force_download=False
                )
                shutil.copy2(downloaded_path, local_path)
                print(f"βœ“ Downloaded and copied {filename}")
            except Exception as e:
                print(f"βœ— Failed to download {filename}: {e}")
                raise e
        else:
            print(f"βœ“ {filename} already exists locally")
    print("All model files are ready!")

def load_model_on_gpu():
    """Load model inside GPU context - only called within @spaces.GPU decorated function."""
    global model, model_loaded
    
    if model_loaded and model is not None:
        return True
    
    if not chatterbox_available:
        print("ERROR: Chatterbox TTS library not available")
        return False
    
    try:
        print("Loading model inside GPU context...")
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Loading model on device: {device}")
        
        # Try different loading methods
        try:
            model = ChatterboxTTS.from_local(LOCAL_MODEL_PATH, device)
            print("βœ“ Model loaded successfully using from_local method.")
        except Exception as e1:
            print(f"from_local failed: {e1}")
            try:
                model = ChatterboxTTS.from_pretrained(device)
                print("βœ“ Model loaded successfully with from_pretrained.")
            except Exception as e2:
                print(f"from_pretrained failed: {e2}")
                model = load_model_manually(device)
        
        if model and hasattr(model, 'to'):
            model = model.to(device)
        if model and hasattr(model, 'eval'):
            model.eval()
        
        model_loaded = True
        print("βœ“ Model loaded successfully in GPU context")
        return True
        
    except Exception as e:
        print(f"ERROR: Failed to load model in GPU context: {e}")
        traceback.print_exc()
        model = None
        model_loaded = False
        return False

def load_model_manually(device):
    """Manual model loading with proper error handling."""
    import pathlib
    import json
    
    model_path = pathlib.Path(LOCAL_MODEL_PATH)
    print("Manual loading with correct constructor signature...")
    
    s3gen_path = model_path / "s3gen.pt"
    ve_path = model_path / "ve.pt"
    tokenizer_path = model_path / "tokenizer.json"
    t3_cfg_path = model_path / "t3_cfg.pt"
    
    s3gen = torch.load(s3gen_path, map_location='cpu')
    ve = torch.load(ve_path, map_location='cpu')
    t3_cfg = torch.load(t3_cfg_path, map_location='cpu')
    
    with open(tokenizer_path, 'r') as f:
        tokenizer_data = json.load(f)
    
    try:
        from chatterbox.models.tokenizers.tokenizer import EnTokenizer
        tokenizer = EnTokenizer.from_dict(tokenizer_data)
    except Exception:
        tokenizer = tokenizer_data
    
    model = ChatterboxTTS(
        t3=t3_cfg,
        s3gen=s3gen,
        ve=ve,
        tokenizer=tokenizer,
        device=device
    )
    
    print("βœ“ Model loaded successfully with manual constructor.")
    return model

def cleanup_gpu_memory():
    """Clean up GPU memory - only call within GPU context."""
    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            gc.collect()
    except Exception as e:
        print(f"Warning: GPU cleanup failed: {e}")

def truncate_text_safely(text, max_chars=MAX_CHARS_PER_GENERATION):
    """Truncate text to safe length while preserving sentence boundaries."""
    if len(text) <= max_chars:
        return text, False
    
    # Find the last sentence ending before the limit
    truncated = text[:max_chars]
    
    # Look for sentence endings
    for ending in ['. ', '! ', '? ']:
        last_sentence = truncated.rfind(ending)
        if last_sentence > max_chars * 0.7:  # Don't truncate too aggressively
            return text[:last_sentence + 1].strip(), True
    
    # Fallback to word boundary
    last_space = truncated.rfind(' ')
    if last_space > max_chars * 0.8:
        return text[:last_space].strip(), True
    
    # Last resort: hard truncate
    return truncated.strip(), True

# Download model files during startup (CPU only)
if chatterbox_available:
    try:
        download_model_files()
        print("Model files downloaded. Model will be loaded on first GPU request.")
    except Exception as e:
        print(f"ERROR during model file download: {e}")

@spaces.GPU
def clone_voice(text_to_speak, reference_audio_path, exaggeration=0.6, cfg_pace=0.3, random_seed=0, temperature=0.6):
    """Main voice cloning function - runs on GPU."""
    global model, model_loaded
    
    # Input validation
    if not chatterbox_available:
        return None, "Error: Chatterbox TTS library not available. Please check installation."
    
    if not text_to_speak or text_to_speak.strip() == "":
        return None, "Error: Please enter some text to speak."
    
    if reference_audio_path is None:
        return None, "Error: Please upload a reference audio file (.wav or .mp3)."

    # Check text length and truncate if necessary
    original_length = len(text_to_speak)
    if original_length > MAX_CHARS_TOTAL:
        return None, f"Error: Text is too long ({original_length:,} characters). Maximum allowed is {MAX_CHARS_TOTAL:,} characters. Please use the chunked generation API for longer texts."
    
    # Truncate to safe generation length
    text_to_use, was_truncated = truncate_text_safely(text_to_speak, MAX_CHARS_PER_GENERATION)
    
    try:
        # Load model if not already loaded
        if not model_loaded:
            print("Loading model for the first time...")
            if not load_model_on_gpu():
                return None, "Error: Failed to load model. Please check the logs for details."
        
        if model is None:
            return None, "Error: Model not loaded. Please check the logs for details."
        
        print(f"Processing request:")
        print(f"  Original text length: {original_length:,} characters")
        print(f"  Processing length: {len(text_to_use):,} characters")
        print(f"  Truncated: {was_truncated}")
        print(f"  Audio: '{reference_audio_path}'")
        print(f"  Parameters: exag={exaggeration}, cfg={cfg_pace}, seed={random_seed}, temp={temperature}")
        
        # Clean GPU memory before generation
        cleanup_gpu_memory()
        
        # Set random seed if specified
        if random_seed > 0:
            torch.manual_seed(random_seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(random_seed)
        
        # Check CUDA availability and memory
        if torch.cuda.is_available():
            print(f"CUDA memory before generation: {torch.cuda.memory_allocated() / 1024**2:.1f} MB")
        
        # Generate audio with error handling
        try:
            with torch.no_grad():
                output_wav_data = model.generate(
                    text=text_to_use,
                    audio_prompt_path=reference_audio_path,
                    exaggeration=exaggeration,
                    cfg_weight=cfg_pace,
                    temperature=temperature
                )
        except RuntimeError as e:
            if "CUDA" in str(e) or "out of memory" in str(e) or "device-side assert" in str(e):
                print(f"CUDA error during generation: {e}")
                cleanup_gpu_memory()
                return None, f"CUDA error: Text may be too long for single generation. Try shorter text (under {MAX_CHARS_PER_GENERATION} characters) or use the chunked generation API for longer content."
            else:
                raise e
        
        # Get sample rate
        try:
            sample_rate = model.sr
        except:
            sample_rate = 24000
        
        # Process output
        if isinstance(output_wav_data, str):
            result = output_wav_data
        else:
            import numpy as np
            if hasattr(output_wav_data, 'cpu'):
                output_wav_data = output_wav_data.cpu().numpy()
            if output_wav_data.ndim > 1:
                output_wav_data = output_wav_data.squeeze()
            result = (sample_rate, output_wav_data)
        
        # Clean up GPU memory after generation
        cleanup_gpu_memory()
        
        if torch.cuda.is_available():
            print(f"CUDA memory after generation: {torch.cuda.memory_allocated() / 1024**2:.1f} MB")
        
        print("βœ“ Audio generated successfully")
        
        # Prepare success message
        success_msg = "Success: Audio generated successfully!"
        if was_truncated:
            success_msg += f" Note: Text was truncated from {original_length:,} to {len(text_to_use):,} characters for optimal generation. Use the chunked generation API for longer texts."
        
        return result, success_msg
        
    except Exception as e:
        print(f"ERROR during audio generation: {e}")
        traceback.print_exc()
        
        # Clean up on error
        try:
            cleanup_gpu_memory()
        except:
            pass
        
        # Provide specific error messages
        error_msg = str(e)
        if "CUDA" in error_msg or "device-side assert" in error_msg:
            return None, f"CUDA error: {error_msg}. Try shorter text (under {MAX_CHARS_PER_GENERATION} characters) or use the chunked generation API."
        elif "out of memory" in error_msg:
            return None, f"GPU memory error: {error_msg}. Please try with shorter text."
        else:
            return None, f"Error during audio generation: {error_msg}. Check logs for more details."

def clone_voice_api(text_to_speak, reference_audio_url, exaggeration=0.6, cfg_pace=0.3, random_seed=0, temperature=0.6):
    """API wrapper function."""
    import requests
    import tempfile
    import os
    import base64

    temp_audio_path = None
    try:
        # Handle different audio input formats
        if reference_audio_url.startswith('data:audio'):
            header, encoded = reference_audio_url.split(',', 1)
            audio_data = base64.b64decode(encoded)
            ext = '.mp3' if 'mp3' in header else '.wav'
            with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:
                temp_file.write(audio_data)
                temp_audio_path = temp_file.name
        elif reference_audio_url.startswith('http'):
            response = requests.get(reference_audio_url, timeout=30)
            response.raise_for_status()
            ext = '.mp3' if reference_audio_url.endswith('.mp3') else '.wav'
            with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:
                temp_file.write(response.content)
                temp_audio_path = temp_file.name
        else:
            temp_audio_path = reference_audio_url

        # Call the GPU function
        audio_output, status = clone_voice(text_to_speak, temp_audio_path, exaggeration, cfg_pace, random_seed, temperature)
        
        return audio_output, status
        
    except Exception as e:
        print(f"API Error: {e}")
        return None, f"API Error: {str(e)}"
    finally:
        # Clean up temporary file
        if temp_audio_path and temp_audio_path != reference_audio_url:
            try:
                os.unlink(temp_audio_path)
            except:
                pass

def main():
    print("Starting Advanced Gradio interface...")
    
    with gr.Blocks(title="πŸŽ™οΈ Advanced Chatterbox Voice Cloning") as demo:
        gr.Markdown("# πŸŽ™οΈ Advanced Chatterbox Voice Cloning")
        gr.Markdown("Clone any voice using advanced AI technology with fine-tuned controls.")
        
        # Add warning about text length
        gr.Markdown(f"""
        **⚠️ Text Length Limits:**
        - **Single Generation**: Up to {MAX_CHARS_PER_GENERATION:,} characters (optimal quality)
        - **API Maximum**: Up to {MAX_CHARS_TOTAL:,} characters (may be truncated)
        - **For longer texts**: Use the chunked generation API in your application
        """)

        with gr.Row():
            with gr.Column(scale=2):
                text_input = gr.Textbox(
                    label=f"Text to Speak (max {MAX_CHARS_TOTAL:,} characters)",
                    placeholder="Enter the text you want the cloned voice to say...",
                    lines=5,
                    max_lines=10
                )
                audio_input = gr.Audio(
                    type="filepath",
                    label="Reference Audio (Upload a short .wav or .mp3 clip)",
                    sources=["upload", "microphone"]
                )

                with gr.Accordion("πŸ”§ Advanced Settings", open=False):
                    with gr.Row():
                        exaggeration_input = gr.Slider(
                            minimum=0.25, maximum=1.0, value=0.6, step=0.05,
                            label="Exaggeration", info="Controls voice characteristic emphasis"
                        )
                        cfg_pace_input = gr.Slider(
                            minimum=0.2, maximum=1.0, value=0.3, step=0.05,
                            label="CFG/Pace", info="Classifier-free guidance weight"
                        )
                    with gr.Row():
                        seed_input = gr.Number(
                            value=0, label="Random Seed", info="Set to 0 for random results", precision=0
                        )
                        temperature_input = gr.Slider(
                            minimum=0.05, maximum=2.0, value=0.6, step=0.05,
                            label="Temperature", info="Controls randomness in generation"
                        )

                generate_btn = gr.Button("🎡 Generate Voice Clone", variant="primary", size="lg")

            with gr.Column(scale=1):
                audio_output = gr.Audio(label="Generated Audio", type="numpy")
                status_output = gr.Textbox(label="Status", lines=3)

        # Connect the interface
        generate_btn.click(
            fn=clone_voice_api,
            inputs=[text_input, audio_input, exaggeration_input, cfg_pace_input, seed_input, temperature_input],
            outputs=[audio_output, status_output],
            api_name="predict"
        )

        # API endpoint for external calls
        def clone_voice_base64_api(text_to_speak, reference_audio_b64, exaggeration=0.6, cfg_pace=0.3, random_seed=0, temperature=0.6):
            return clone_voice_api(text_to_speak, reference_audio_b64, exaggeration, cfg_pace, random_seed, temperature)

        # Hidden API interface
        with gr.Row(visible=False):
            api_text_input = gr.Textbox()
            api_audio_input = gr.Textbox()
            api_exaggeration_input = gr.Slider(minimum=0.25, maximum=1.0, value=0.6)
            api_cfg_pace_input = gr.Slider(minimum=0.2, maximum=1.0, value=0.3)
            api_seed_input = gr.Number(value=0, precision=0)
            api_temperature_input = gr.Slider(minimum=0.05, maximum=2.0, value=0.6)
            api_audio_output = gr.Audio(type="numpy")
            api_status_output = gr.Textbox()
            api_btn = gr.Button()

        api_btn.click(
            fn=clone_voice_base64_api,
            inputs=[api_text_input, api_audio_input, api_exaggeration_input, api_cfg_pace_input, api_seed_input, api_temperature_input],
            outputs=[api_audio_output, api_status_output],
            api_name="clone_voice"
        )

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        quiet=False,
        share=False
    )

if __name__ == "__main__":
    main()