Create app.py

app.py

@@ -0,0 +1,167 @@
+"""
+Custom captioning tool: transcribes audio/video with Whisper (large-v3) and chunks
+captions by a user-specified number of words per caption, outputting an SRT file
+ready to import into CapCut or any other editor.
+
+Designed for Hugging Face Spaces ZeroGPU (dynamic H200 access).
+Uses openai-whisper (PyTorch-based) rather than faster-whisper/CTranslate2,
+since ZeroGPU's GPU-attach mechanism is built around torch.cuda and is most
+reliable with PyTorch-native models.
+"""
+
+import os
+import tempfile
+
+import gradio as gr
+import spaces
+import torch
+import whisper
+
+# ---------------------------------------------------------------------------
+# Model loads once at startup on CPU. It's moved to GPU only inside the
+# @spaces.GPU-decorated function below, since real CUDA access on ZeroGPU
+# only exists during that call.
+# ---------------------------------------------------------------------------
+MODEL_SIZE = "large-v3"
+model = whisper.load_model(MODEL_SIZE, device="cpu")
+
+
+def format_timestamp(seconds: float) -> str:
+    """Convert seconds (float) to SRT timestamp format: HH:MM:SS,mmm"""
+    ms_total = int(round(seconds * 1000))
+    hours, rem = divmod(ms_total, 3600_000)
+    minutes, rem = divmod(rem, 60_000)
+    secs, ms = divmod(rem, 1000)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"
+
+
+def chunk_words(words, words_per_caption, max_chars=None):
+    """
+    Group a flat list of {"word","start","end"} dicts into caption chunks
+    of `words_per_caption` words each. Optionally caps chunk length by
+    max_chars so long words don't overflow a caption line.
+    """
+    captions = []
+    current = []
+
+    def flush():
+        if current:
+            captions.append({
+                "text": " ".join(w["word"] for w in current).strip(),
+                "start": current[0]["start"],
+                "end": current[-1]["end"],
+            })
+
+    for w in words:
+        current.append(w)
+        text_len = len(" ".join(x["word"] for x in current))
+        hit_word_limit = len(current) >= words_per_caption
+        hit_char_limit = max_chars is not None and text_len >= max_chars
+        if hit_word_limit or hit_char_limit:
+            flush()
+            current = []
+    flush()
+
+    return captions
+
+
+def write_srt(captions, path):
+    with open(path, "w", encoding="utf-8") as f:
+        for i, cap in enumerate(captions, start=1):
+            f.write(f"{i}\n")
+            f.write(f"{format_timestamp(cap['start'])} --> {format_timestamp(cap['end'])}\n")
+            f.write(f"{cap['text']}\n\n")
+
+
+@spaces.GPU(duration=120)  # per-call GPU time budget; counts against your daily quota
+def run_transcription(media_path):
+    """The only part that touches CUDA -- kept as small as possible to conserve quota."""
+    global model
+    model = model.to("cuda")
+    result = model.transcribe(media_path, word_timestamps=True, fp16=True)
+    model = model.to("cpu")  # release VRAM before the GPU is handed back
+    torch.cuda.empty_cache()
+    return result
+
+
+def transcribe_and_caption(media_file, words_per_caption, max_chars, progress=gr.Progress()):
+    if media_file is None:
+        return None, "Upload an audio or video file first."
+
+    words_per_caption = max(1, int(words_per_caption))
+    max_chars_val = int(max_chars) if max_chars and max_chars > 0 else None
+
+    progress(0.2, desc="Requesting GPU and transcribing (uses your ZeroGPU quota)...")
+    result = run_transcription(media_file)
+
+    words = []
+    for segment in result.get("segments", []):
+        for w in segment.get("words", []):
+            words.append({
+                "word": w["word"].strip(),
+                "start": w["start"],
+                "end": w["end"],
+            })
+
+    if not words:
+        return None, "No speech detected in the file."
+
+    progress(0.8, desc="Building captions...")
+    captions = chunk_words(words, words_per_caption, max_chars_val)
+
+    out_path = os.path.join(tempfile.gettempdir(), "captions.srt")
+    write_srt(captions, out_path)
+
+    preview_lines = [f"[{format_timestamp(c['start'])}] {c['text']}" for c in captions[:15]]
+    preview = "\n".join(preview_lines)
+    if len(captions) > 15:
+        preview += f"\n... ({len(captions) - 15} more captions)"
+
+    detected_lang = result.get("language", "unknown")
+    summary = (
+        f"Detected language: {detected_lang}\n"
+        f"Total captions: {len(captions)}\n\n"
+        f"Preview:\n{preview}"
+    )
+
+    progress(1.0, desc="Done")
+    return out_path, summary
+
+
+with gr.Blocks(title="Word-Count Caption Generator") as demo:
+    gr.Markdown(
+        "# Word-Count Caption Generator\n"
+        "Upload a video or audio file, choose how many words should appear per caption, "
+        "and get back an `.srt` file to import into CapCut (or any editor).\n\n"
+        f"Running **{MODEL_SIZE}** on ZeroGPU (H200). Each run uses part of your daily "
+        "GPU quota, so batch your clips rather than testing repeatedly."
+    )
+
+    with gr.Row():
+        with gr.Column():
+            media_input = gr.File(
+                label="Audio or video file",
+                file_types=["audio", "video"],
+            )
+            words_slider = gr.Slider(
+                minimum=1, maximum=10, value=3, step=1,
+                label="Words per caption",
+            )
+            chars_slider = gr.Slider(
+                minimum=0, maximum=60, value=0, step=1,
+                label="Max characters per caption (0 = no limit)",
+            )
+            run_btn = gr.Button("Generate captions", variant="primary")
+
+        with gr.Column():
+            srt_output = gr.File(label="Download .srt")
+            summary_output = gr.Textbox(label="Summary / preview", lines=18)
+
+    run_btn.click(
+        fn=transcribe_and_caption,
+        inputs=[media_input, words_slider, chars_slider],
+        outputs=[srt_output, summary_output],
+    )
+
+if __name__ == "__main__":
+    demo.launch()