Update app.py

app.py

@@ -20,6 +20,7 @@ generate_kwargs = {
     "num_beams": 1,
     "no_repeat_ngram_size": 5,
     "max_new_tokens": 64,
+    "return_timestamps": True,  # Necesario para obtener los tiempos
 }
 
 model_dict = {
@@ -41,6 +42,18 @@ pipe_dict = {
 logger.success("Pipelines initialized!")
 
 
+def save_as_srt(transcription, timestamps, output_path):
+    """Genera un archivo .srt a partir de las transcripciones y sus marcas de tiempo."""
+    with open(output_path, "w", encoding="utf-8") as f:
+        for idx, (text, (start, end)) in enumerate(zip(transcription, timestamps)):
+            start_time = time.strftime('%H:%M:%S', time.gmtime(start)) + f",{int(start % 1 * 1000):03d}"
+            end_time = time.strftime('%H:%M:%S', time.gmtime(end)) + f",{int(end % 1 * 1000):03d}"
+
+            f.write(f"{idx + 1}\n")
+            f.write(f"{start_time} --> {end_time}\n")
+            f.write(f"{text}\n\n")
+
+
 @spaces.GPU
 def transcribe_common(audio: str, model: str) -> str:
     if not audio:
@@ -48,7 +61,7 @@ def transcribe_common(audio: str, model: str) -> str:
     filename = Path(audio).name
     logger.info(f"Model: {model}")
     logger.info(f"Audio: {filename}")
-    
+
     try:
         y, sr = librosa.load(audio, mono=True, sr=16000)
     except Exception as e:
@@ -60,22 +73,22 @@ def transcribe_common(audio: str, model: str) -> str:
 
     duration = librosa.get_duration(y=y, sr=sr)
     logger.info(f"Duration: {duration:.2f}s")
-    kwargs = generate_kwargs.copy()
-    if duration > 30:
-        kwargs["return_timestamps"] = True
 
     start_time = time.time()
-    result = pipe_dict[model](y, generate_kwargs=kwargs)["text"]
+    result = pipe_dict[model](y, generate_kwargs=generate_kwargs)
     end_time = time.time()
-    logger.success(f"Finished in {end_time - start_time:.2f}s\n{result}")
-    
-    # Guardar resultado en un archivo .str
-    output_path = f"{Path(filename).stem}_{model}.str"
-    with open(output_path, "w", encoding="utf-8") as f:
-        f.write(result)
+
+    transcription = result["text"]
+    timestamps = result["chunks"]  # Esto contiene las marcas de tiempo
     
+    logger.success(f"Finished in {end_time - start_time:.2f}s\n{transcription}")
+
+    # Guardar resultado en un archivo .srt
+    output_path = f"{Path(filename).stem}.srt"
+    save_as_srt([chunk["text"] for chunk in timestamps], [(chunk["timestamp_start"], chunk["timestamp_end"]) for chunk in timestamps], output_path)
+
     logger.info(f"Transcription saved to {output_path}")
-    return output_path  # Devuelve el path del archivo transcrito
+    return transcription
 
 
 def transcribe_others(audio) -> tuple[str, str]:
@@ -90,7 +103,7 @@ def transcribe_anime_whisper(audio) -> str:
 
 initial_md = """
 # Anime-Whisper Demo
-[**Anime Whisper**](https://huggingface.co/litagin/anime-whisper): 5千時間以上のアニメ調セリフと台本でファインチューニングされた日本語音声認識モデルのデモです。句読点や感嘆符がリズムや感情に合わせて自然に付き、NSFW含む非言語発言もうまく台本調に書き起こされます。
+[**Anime Whisper**](https://huggingface.co/litagin/anime-whisper): 5千時間以上のアニメ調セリフと台本でファインチューニングされた日本語音声認識モデルのデモです。句読点や感嘆符がリズムや感情に合わせて自然に付き、NSFW含む非言語発話もうまく台本調に書き起こされます。
 - デモでは**音声は15秒まで**しか受け付けません
 - 日本語のみ対応 (Japanese only)
 - 比較のために [openai/whisper-large-v3-turbo](https://huggingface.co/openai/whisper-large-v3-turbo) と [kotoba-tech/kotoba-whisper-v2.0](https://huggingface.co/kotoba-tech/kotoba-whisper-v2.0) も用意しています
@@ -102,6 +115,7 @@ generate_kwargs = {
     "num_beams": 1,
     "no_repeat_ngram_size": 5,
     "max_new_tokens": 64,  # 結果が長いときは途中で打ち切られる
+    "return_timestamps": True,  # Para incluir tiempos
 }
 ```
 """
@@ -113,16 +127,16 @@ with gr.Blocks() as app:
         with gr.Column():
             gr.Markdown("### Anime-Whisper")
             button_galgame = gr.Button("Transcribe with Anime-Whisper")
-            output_galgame = gr.File(label="Download Anime-Whisper Transcription")
+            output_galgame = gr.Textbox(label="Result")
     gr.Markdown("### Comparison")
     button_others = gr.Button("Transcribe with other models")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Whisper-Large-V3-Turbo")
-            output_v3 = gr.File(label="Download Whisper-Large-V3-Turbo Transcription")
+            output_v3 = gr.Textbox(label="Result")
         with gr.Column():
             gr.Markdown("### Kotoba-Whisper-V2.0")
-            output_kotoba_v2 = gr.File(label="Download Kotoba-Whisper-V2.0 Transcription")
+            output_kotoba_v2 = gr.Textbox(label="Result")
 
     button_galgame.click(
         transcribe_anime_whisper,