Update app.py

app.py

@@ -1,25 +1,36 @@
-import spaces
-from kokoro import KModel, KPipeline
 import gradio as gr
-import os
+import openai
+from kokoro import KPipeline
 import random
+import os
 import torch
+import time
 
-IS_DUPLICATE = not os.getenv('SPACE_ID', '').startswith('hexgrad/')
-CHAR_LIMIT = None if IS_DUPLICATE else 5000
+# Set up the OpenAI API key (optional)
+openai.api_key = None  # Will be set by the user through the UI
 
+# Check if GPU is available
 CUDA_AVAILABLE = torch.cuda.is_available()
+
+# Initialize the models and pipelines (for TTS)
 models = {gpu: KModel().to('cuda' if gpu else 'cpu').eval() for gpu in [False] + ([True] if CUDA_AVAILABLE else [])}
-pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'abefhijpz'}
+pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'abefhijpz']
+# Load lexicon for specific languages
 pipelines['a'].g2p.lexicon.golds['kokoro'] = 'kˈOkəɹO'
 pipelines['b'].g2p.lexicon.golds['kokoro'] = 'kˈQkəɹQ'
 
-@spaces.GPU(duration=10)
-def forward_gpu(ps, ref_s, speed):
-    return models[True](ps, ref_s, speed)
+# Initialize random texts for generating sample text
+random_texts = {}
+for lang in ['en']:
+    with open(f'{lang}.txt', 'r') as r:
+        random_texts[lang] = [line.strip() for line in r]
+
+def get_random_text(voice):
+    lang = dict(a='en', b='en')[voice[0]]
+    return random.choice(random_texts[lang])
 
+# Generate function to create speech from text
 def generate_first(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
-    text = text if CHAR_LIMIT is None else text.strip()[:CHAR_LIMIT]
     pipeline = pipelines[voice[0]]
     pack = pipeline.load_voice(voice)
     use_gpu = use_gpu and CUDA_AVAILABLE
@@ -40,18 +51,23 @@ def generate_first(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
         return (24000, audio.numpy()), ps
     return None, ''
 
-# Arena API
-def predict(text, voice='af_heart', speed=1):
-    return generate_first(text, voice, speed, use_gpu=False)[0]
-
-def tokenize_first(text, voice='af_heart'):
-    # Split the input text into words and return as a list of words (fix applied here)
-    words = text.split()  # This splits the text into words based on spaces
-    return words  # Return a list of words
+# Translator function using OpenAI API
+def translate_to_english(api_key, text, lang_code):
+    openai.api_key = api_key
+    try:
+        prompt = f"Translate the following text from {lang_code} to English: \n\n{text}"
+        response = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages=[{"role": "system", "content": "You are a helpful assistant that translates text."},
+                      {"role": "user", "content": prompt}]
+        )
+        translated_text = response['choices'][0]['message']['content'].strip()
+        return translated_text
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-def generate_all(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
-    text = text if CHAR_LIMIT is None else text.strip()[:CHAR_LIMIT]
-    pipeline = pipelines[voice[0]]
+def generate_audio_from_text(text, lang_code, voice, speed, use_gpu=True):
+    pipeline = pipelines[lang_code]
     pack = pipeline.load_voice(voice)
     use_gpu = use_gpu and CUDA_AVAILABLE
     for _, ps, _ in pipeline(text, voice, speed):
@@ -68,131 +84,114 @@ def generate_all(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
                 audio = models[False](ps, ref_s, speed)
             else:
                 raise gr.Error(e)
-        yield 24000, audio.numpy()
-
-random_texts = {}
-for lang in ['en']:
-    with open(f'{lang}.txt', 'r') as r:
-        random_texts[lang] = [line.strip() for line in r]
-
-def get_random_text(voice):
-    lang = dict(a='en', b='en')[voice[0]]
-    return random.choice(random_texts[lang])
+        return (24000, audio.numpy())
 
+# Define your available voices here in the CHOICES dictionary
 CHOICES = {
-'🇺🇸 🚺 Heart ❤️': 'af_heart',
-'🇺🇸 🚺 Bella 🔥': 'af_bella',
-'🇺🇸 🚺 Nicole 🎧': 'af_nicole',
-'🇺🇸 🚺 Aoede': 'af_aoede',
-'🇺🇸 🚺 Kore': 'af_kore',
-'🇺🇸 🚺 Sarah': 'af_sarah',
-'🇺🇸 🚺 Nova': 'af_nova',
-'🇺🇸 🚺 Sky': 'af_sky',
-'🇺🇸 🚺 Alloy': 'af_alloy',
-'🇺🇸 🚺 Jessica': 'af_jessica',
-'🇺🇸 🚺 River': 'af_river',
-    
-'🇺🇸 🚹 Michael': 'am_michael',
-'🇺🇸 🚹 Fenrir': 'am_fenrir',
-'🇺🇸 🚹 Puck': 'am_puck',
-'🇺🇸 🚹 Echo': 'am_echo',
-'🇺🇸 🚹 Eric': 'am_eric',
-'🇺🇸 🚹 Liam': 'am_liam',
-'🇺🇸 🚹 Onyx': 'am_onyx',
-'🇺🇸 🚹 Santa': 'am_santa',
-'🇺🇸 🚹 Adam': 'am_adam',
-    
-'🇬🇧 🚺 Emma': 'bf_emma',
-'🇬🇧 🚺 Isabella': 'bf_isabella',
-'🇬🇧 🚺 Alice': 'bf_alice',
-'🇬🇧 🚺 Lily': 'bf_lily',
-    
-'🇬🇧 🚹 George': 'bm_george',
-'🇬🇧 🚹 Fable': 'bm_fable',
-'🇬🇧 🚹 Lewis': 'bm_lewis',
-'🇬🇧 🚹 Daniel': 'bm_daniel',
-    
-'🇪🇸 🚺 Dora': 'ef_dora',
-    
-'🇪🇸 🚹 Alex': 'em_alex',
-'🇪🇸 🚹 Santa': 'em_santa',
-    
-'🇫🇷 🚺 Siwis': 'ff_siwis',
-    
-'🇮🇳 🚹 Alpha': 'hf_alpha',
-'🇮🇳 🚹 Beta': 'hf_beta',
-    
-'🇮🇳 🚹 Omega': 'hm_omega',
-'🇮🇳 🚹 Psi': 'hm_psi',
-    
-'🇮🇹 🚺 Sara': 'if_sara',
-    
-'🇮🇹 🚺 Nicola': 'im_nicola',
-    
-'🇯🇵 🚹 Alpha': 'jf_alpha',
-'🇯🇵 🚹 Gongitsune': 'jf_gongitsune',
-'🇯🇵 🚹 Nezumi': 'jf_nezumi',
-'🇯🇵 🚹 Tebukuro': 'jf_tebukuro',
-    
-'🇯🇵 🚹 Kumo': 'jm_kumo',
-    
-'🇧🇷 🚺 Dora': 'pf_dora',
-    
-'🇧🇷 🚹 Alex': 'pm_alex',
-'🇧🇷 🚹 Santa': 'pm_santa',
-    
-'🇨🇳 🚺 Xiaobei': 'zf_xiaobei',
-'🇨🇳 🚺 Xiaoni': 'zf_xiaoni',
-'🇨🇳 🚺 Xiaoxiao': 'zf_xiaoxiao',
-'🇨🇳 🚺 Xiaoyi': 'zf_xiaoyi',
-
-'🇨🇳 🚹 Yunjian': 'zm_yunjian',
-'🇨🇳 🚹 Yunxi': 'zm_yunxi',
-'🇨🇳 🚹 Yunxia': 'zm_yunxia',
-'🇨🇳 🚹 Yunyang': 'zm_yunyang',
+    'af_heart': '🇺🇸 🚺 Heart ❤️',
+    'af_bella': '🇺🇸 🚺 Bella 🔥',
+    'af_nicole': '🇺🇸 🚺 Nicole 🎧',
+    'af_aoede': '🇺🇸 🚺 Aoede',
+    'af_kore': '🇺🇸 🚺 Kore',
+    'af_sarah': '🇺🇸 🚺 Sarah',
+    'af_nova': '🇺🇸 🚺 Nova',
+    'af_sky': '🇺🇸 🚺 Sky',
+    'af_alloy': '🇺🇸 🚺 Alloy',
+    'af_jessica': '🇺🇸 🚺 Jessica',
+    'af_river': '🇺🇸 🚺 River',
+    
+    'am_michael': '🇺🇸 🚹 Michael',
+    'am_fenrir': '🇺🇸 🚹 Fenrir',
+    'am_puck': '🇺🇸 🚹 Puck',
+    'am_echo': '🇺🇸 🚹 Echo',
+    'am_eric': '🇺🇸 🚹 Eric',
+    'am_liam': '🇺🇸 🚹 Liam',
+    'am_onyx': '🇺🇸 🚹 Onyx',
+    'am_santa': '🇺🇸 🚹 Santa',
+    'am_adam': '🇺🇸 🚹 Adam',
+    
+    'bf_emma': '🇬🇧 🚺 Emma',
+    'bf_isabella': '🇬🇧 🚺 Isabella',
+    'bf_alice': '🇬🇧 🚺 Alice',
+    'bf_lily': '🇬🇧 🚺 Lily',
+    
+    'bm_george': '🇬🇧 🚹 George',
+    'bm_fable': '🇬🇧 🚹 Fable',
+    'bm_lewis': '🇬🇧 🚹 Lewis',
+    'bm_daniel': '🇬🇧 🚹 Daniel',
+    
+    'ef_dora': '🇪🇸 🚺 Dora',
+    
+    'em_alex': '🇪🇸 🚹 Alex',
+    'em_santa': '🇪🇸 🚹 Santa',
+    
+    'ff_siwis': '🇫🇷 🚺 Siwis',
+    
+    'hf_alpha': '🇮🇳 🚹 Alpha',
+    'hf_beta': '🇮🇳 🚹 Beta',
+    
+    'hm_omega': '🇮🇳 🚹 Omega',
+    'hm_psi': '🇮🇳 🚹 Psi',
+    
+    'if_sara': '🇮🇹 🚺 Sara',
+    
+    'im_nicola': '🇮🇹 🚺 Nicola',
+    
+    'jf_alpha': '🇯🇵 🚹 Alpha',
+    'jf_gongitsune': '🇯🇵 🚹 Gongitsune',
+    'jf_nezumi': '🇯🇵 🚹 Nezumi',
+    'jf_tebukuro': '🇯🇵 🚹 Tebukuro',
+    
+    'jm_kumo': '🇯🇵 🚹 Kumo',
+    
+    'pf_dora': '🇧🇷 🚺 Dora',
+    
+    'pm_alex': '🇧🇷 🚹 Alex',
+    'pm_santa': '🇧🇷 🚹 Santa',
+    
+    'zf_xiaobei': '🇨🇳 🚺 Xiaobei',
+    'zf_xiaoni': '🇨🇳 🚺 Xiaoni',
+    'zf_xiaoxiao': '🇨🇳 🚺 Xiaoxiao',
+    'zf_xiaoyi': '🇨🇳 🚺 Xiaoyi',
+    
+    'zm_yunjian': '🇨🇳 🚹 Yunjian',
+    'zm_yunxi': '🇨🇳 🚹 Yunxi',
+    'zm_yunxia': '🇨🇳 🚹 Yunxia',
+    'zm_yunyang': '🇨🇳 🚹 Yunyang'
 }
-for v in CHOICES.values():
-    pipelines[v[0]].load_voice(v)
-
-with gr.Blocks() as generate_tab:
-    out_audio = gr.Audio(label='Output Audio', interactive=False, streaming=False, autoplay=True)
-    generate_btn = gr.Button('Generate', variant='primary')
-    with gr.Accordion('Output Tokens', open=True):
-        out_ps = gr.Textbox(interactive=False, show_label=False, info='Tokens used to generate the audio, up to 510 context length.')
-        tokenize_btn = gr.Button('Tokenize', variant='secondary')
-        predict_btn = gr.Button('Predict', variant='secondary', visible=False)
-
-BANNER_TEXT = '''
-[***Kokoro*** **is an open-weight TTS model with 82 million parameters.**](https://huggingface.co/hexgrad/Kokoro-82M)
-As of January 31st, 2025, Kokoro was the most-liked [**TTS model**](https://huggingface.co/models?pipeline_tag=text-to-speech&sort=likes) and the most-liked [**TTS space**](https://huggingface.co/spaces?sort=likes&search=tts) on Hugging Face.
-This demo only showcases English, but you can directly use the model to access other languages.
-'''
-
-API_OPEN = os.getenv('SPACE_ID') != 'hexgrad/Kokoro-TTS'
-API_NAME = None if API_OPEN else False
+
+# Gradio interface setup
 with gr.Blocks() as app:
-    with gr.Row():
-        gr.Markdown(BANNER_TEXT, container=True)
+    gr.Markdown("### Kokoro Text-to-Speech with Translation")
     with gr.Row():
         with gr.Column():
-            text = gr.Textbox(label='Input Text', info=f"Up to ~500 characters per Generate, or {'∞' if CHAR_LIMIT is None else CHAR_LIMIT} characters per Stream")
-            with gr.Row():
-                voice = gr.Dropdown(list(CHOICES.items()), value='af_heart', label='Voice', info='Quality and availability vary by language')
-                use_gpu = gr.Dropdown(
-                    [('ZeroGPU 🚀', True), ('CPU 🐌', False)],
-                    value=CUDA_AVAILABLE,
-                    label='Hardware',
-                    info='GPU is usually faster, but has a usage quota',
-                    interactive=CUDA_AVAILABLE
-                )
-            speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='Speed')
-            random_btn = gr.Button('Random Text', variant='secondary')
+            # Input for text and language settings
+            input_text = gr.Textbox(label="Enter Text", placeholder="Type your text here...")
+            voice = gr.Dropdown(list(CHOICES.items()), value='af_heart', label='Voice')
+            use_gpu = gr.Checkbox(label="Use GPU", value=CUDA_AVAILABLE)
+            speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label="Speed")
+            openai_api_key = gr.Textbox(label="Enter OpenAI API Key (for translation)", type="password")
+            random_btn = gr.Button("Random Text")
+
         with gr.Column():
-            gr.TabbedInterface([generate_tab], ['Generate'])
-    random_btn.click(fn=get_random_text, inputs=[voice], outputs=[text], api_name=API_NAME)
-    generate_btn.click(fn=generate_first, inputs=[text, voice, speed, use_gpu], outputs=[out_audio, out_ps], api_name=API_NAME)
-    tokenize_btn.click(fn=tokenize_first, inputs=[text, voice], outputs=[out_ps], api_name=API_NAME)
-    predict_btn.click(fn=predict, inputs=[text, voice, speed], outputs=[out_audio], api_name=API_NAME)
-
-if __name__ == '__main__':
-    app.queue(api_open=API_OPEN).launch(show_api=API_OPEN, ssr_mode=True)
+            out_audio = gr.Audio(label="Generated Audio", interactive=False, autoplay=True)
+            out_text = gr.Textbox(label="Generated Audio Tokens", interactive=False)
+            generate_btn = gr.Button("Generate Audio")
+            translate_btn = gr.Button("Translate and Generate Audio")
+
+        random_btn.click(fn=get_random_text, inputs=[voice], outputs=[input_text])
+        
+        def handle_translation(text, api_key, lang_code, voice, speed, use_gpu):
+            translated_text = translate_to_english(api_key, text, lang_code)
+            translated_audio = generate_audio_from_text(translated_text, 'a', voice, speed, use_gpu)
+            return translated_audio, translated_text
+
+        translate_btn.click(fn=handle_translation, inputs=[input_text, openai_api_key, voice, speed, use_gpu], outputs=[out_audio, out_text])
+
+        def generate_and_play(text, voice, speed, use_gpu):
+            audio, tokens = generate_first(text, voice, speed, use_gpu)
+            return audio, tokens
+
+        generate_btn.click(fn=generate_and_play, inputs=[input_text, voice, speed, use_gpu], outputs=[out_audio, out_text])
+
+    app.launch()