Update app.py

app.py

@@ -9,6 +9,11 @@ import logging
 import torch
 import torchaudio
 from fastapi import FastAPI, Request, HTTPException, BackgroundTasks
+from fastapi.middleware.cors import CORSMiddleware
+from faster_whisper import WhisperModel
+from df.enhance import enhance, init_df, load_audio, save_audio
+from pyannote.audio import Pipeline as DiarizationPipeline
+from pipeline.postprocess import process
 
 
 class HealthAndStatusFilter(logging.Filter):
@@ -19,11 +24,6 @@ class HealthAndStatusFilter(logging.Filter):
 
 
 logging.getLogger("uvicorn.access").addFilter(HealthAndStatusFilter())
-from fastapi.middleware.cors import CORSMiddleware
-from faster_whisper import WhisperModel
-from df.enhance import enhance, init_df, load_audio, save_audio
-from pyannote.audio import Pipeline as DiarizationPipeline
-from pipeline.postprocess import process
 
 torch.set_num_threads(2)  # جلوی رقابت بی‌فایده‌ی رشته‌ها روی CPU محدود رو می‌گیره
 
@@ -62,14 +62,14 @@ print("Speaker-diarization model ready!")
 
 MIN_NOISE_BLOCK_SEC = 1.0
 JOB_TTL_SECONDS = 3 * 60 * 60  # کارهای قدیمی بعد از ۳ ساعت پاک میشن (برای فایل‌های سنگین که ممکنه پردازششون طولانی بشه)
-STAGE_TIMEOUT_SEC = 1200  # سقف زمانی برای هر مرحله‌ی سنگین (۲۰ دقیقه - فقط محافظ آخر در برابر گیرکردن واقعی، نه محدودیت عادی)
+STAGE_TIMEOUT_SEC = 1200  # سقف زمانی برای هر مرحله‌ی سنگین (۲۰ دقیقه - فقط محافظ آخر در برابر گیرکردن واقعی)
 
 # ─── نگهداری وضعیت کارها (Jobs) توی حافظه ───────────────────────────────────
 jobs = {}
 jobs_lock = threading.Lock()
 
 
-def new_job():
+def new_job(single_speaker_mode=False):
     job_id = uuid.uuid4().hex
     tmp_dir = tempfile.mkdtemp(prefix=f"subflow_{job_id}_")
     with jobs_lock:
@@ -83,6 +83,7 @@ def new_job():
             "language": None,
             "duration": None,
             "error": None,
+            "single_speaker_mode": single_speaker_mode,
         }
     return job_id
 
@@ -113,7 +114,7 @@ def cleanup_expired_jobs():
         print(f"[cleanup] removed expired job {jid}")
 
 
-# ─── منطق تشخیص گوینده و هم‌ترازی (بدون تغییر نسبت به نسخه‌ی قبلی) ──────────
+# ─── منطق تشخیص گوینده و هم‌ترازی ────────────────────────────────────────────
 
 def build_speaker_lookup(diarization):
     """
@@ -248,6 +249,7 @@ def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end, cut
         print(f"  [align] skipped due to: {type(e).__name__}: {e}")
         return words
 
+
 def run_transcription(audio_path):
     """
     نکته‌ی مهم: model.transcribe فوراً برنمی‌گرده - یه نتیجه‌ی "تنبل" میده که کار
@@ -269,52 +271,11 @@ def run_transcription(audio_path):
     return segments_list, info
 
 
-def probe_speaker_count(audio_path, job_id=None):
-    """
-    قبل از اجرای تشخیص گوینده‌ی کامل (که سنگین و کنده)، یه آزمایش سریع انجام می‌ده:
-    از ۳ نقطه‌ی پخش‌شده در طول فایل (نه فقط اول)، هرکدوم یه تیکه‌ی ۲۰ ثانیه‌ای برمی‌داره
-    و روی همون تیکه‌های کوتاه تشخیص گوینده اجرا می‌کنه.
-    اگه هر سه فقط یه گوینده نشون بدن، برمی‌گردونه True (یعنی «به‌احتمال زیاد تک‌گوینده‌ست»)
-    وگرنه False (یعنی «باید تشخیص کامل روی کل فایل اجرا بشه»).
-    """
-    info = torchaudio.info(audio_path)
-    duration = info.num_frames / info.sample_rate
-    probe_len = 20.0
-
-    if duration <= probe_len * 3:
-        # فایل خیلی کوتاهه، آزمایش جداگانه فایده‌ای نداره - مستقیم برو سراغ تشخیص کامل
-        return False
-
-    fractions = (0.1, 0.5, 0.85)
-    waveform, sr = torchaudio.load(audio_path)
-
-    for frac in fractions:
-        center = duration * frac
-        start_sec = max(0, center - probe_len / 2)
-        end_sec = min(duration, start_sec + probe_len)
-        start_sample = int(start_sec * sr)
-        end_sample = int(end_sec * sr)
-        clip = waveform[:, start_sample:end_sample]
-
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            torchaudio.save(tmp.name, clip, sr)
-            clip_path = tmp.name
-
-        try:
-            probe_result = diarize_pipeline(clip_path, max_speakers=8)
-            speakers_found = {speaker for _, _, speaker in probe_result.itertracks(yield_label=True)}
-            print(f"  [speaker-probe] window at {frac:.0%} of file: {len(speakers_found)} speaker(s)")
-            if len(speakers_found) > 1:
-                return False
-        finally:
-            os.remove(clip_path)
-
-    return True
-
 class SingleSpeakerResult:
     """
     یه نتیجه‌ی «قلابی» که همون شکل خروجی diarize_pipeline رو تقلید می‌کنه (متد itertracks داره)
     ولی هیچ پردازش واقعی‌ای روی کل فایل انجام نداده - فقط می‌گه «همه‌جا همون یه گوینده‌ست».
+    فقط وقتی استفاده میشه که خودِ کاربر صراحتاً بگه فایل تک‌گوینده‌ست.
     """
     def __init__(self, duration):
         self.duration = duration
@@ -327,18 +288,13 @@ class SingleSpeakerResult:
         yield FakeTurn(0, self.duration), None, "SPEAKER_00"
 
 
-def run_diarization(audio_path, job_id=None):
-    if job_id is not None:
-        update_job(job_id, detail="بررسی سریع تعداد گوینده‌ها...")
-    single_speaker = probe_speaker_count(audio_path, job_id)
-
-    if single_speaker:
-        print("  [speaker-probe] looks single-speaker across all 3 windows — skipping full diarization entirely")
+def run_diarization(audio_path, single_speaker_mode=False, job_id=None):
+    if single_speaker_mode:
+        print("  [diarization] single-speaker mode selected by user — skipping full diarization entirely")
         info = torchaudio.info(audio_path)
         duration = info.num_frames / info.sample_rate
         return SingleSpeakerResult(duration)
 
-    print("  [speaker-probe] multiple speakers detected — running full search")
     return diarize_pipeline(audio_path, max_speakers=8)
 
 
@@ -377,9 +333,10 @@ class ProcessingTimeout(Exception):
 
 def run_blocking_with_timeout(job_id, blocking_fn, *args, status, detail, max_seconds=STAGE_TIMEOUT_SEC):
     """
-    نسخه‌ی هم‌زمان (نه async) - چون این خودش داخل یه thread پس‌زمینه اجرا میشه،
-    نیازی به هماهنگی با event loop نداره. وضعیت Job رو مستقیم آپدیت می‌کنه
-    تا polling بتونه ببینتش.
+    یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه،
+    هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی
+    قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره.
+    اگه از max_seconds بیشتر طول بکشه، به‌جای قفل ابدی، خطای مشخص میده.
     """
     result_holder = {}
     error_holder = {}
@@ -412,6 +369,7 @@ def process_job(job_id):
     """کل پایپلاین پردازش - این تابع توی یه thread پس‌زمینه اجرا میشه، نه توی درخواست HTTP."""
     job = get_job(job_id)
     tmp_dir = job["tmp_dir"]
+    single_speaker_mode = job.get("single_speaker_mode", False)
     tmp_path = os.path.join(tmp_dir, "input.wav")
     enhanced_path = None
 
@@ -431,23 +389,27 @@ def process_job(job_id):
             enhanced_path = os.path.join(tmp_dir, "enhanced.wav")
             save_audio(enhanced_path, enhanced_audio, df_state.sr())
             audio_to_transcribe = enhanced_path
-        
+
         segments_raw_list, info = run_blocking_with_timeout(
             job_id, run_transcription, audio_to_transcribe,
             status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
         )
         update_job(job_id, detail=f"{len(segments_raw_list)} جمله تشخیص داده شد")
-        
+
         align_waveform, sr = torchaudio.load(audio_to_transcribe)
         if sr != 16000:
             align_waveform = torchaudio.functional.resample(align_waveform, sr, 16000)
         if align_waveform.shape[0] > 1:
             align_waveform = align_waveform.mean(dim=0, keepdim=True)
 
-        diarization = run_blocking_with_timeout(
-            job_id, run_diarization, audio_to_transcribe, job_id,
-            status="aligning", detail="در حال تشخیص گوینده‌ها",
-        )
+        if single_speaker_mode:
+            diarization = run_diarization(audio_to_transcribe, single_speaker_mode=True)
+            update_job(job_id, status="aligning", detail="تشخیص گوینده رد شد (حالت تک‌گوینده)")
+        else:
+            diarization = run_blocking_with_timeout(
+                job_id, run_diarization, audio_to_transcribe, False, job_id,
+                status="aligning", detail="در حال تشخیص گوینده‌ها",
+            )
         update_job(job_id, detail="گوینده‌ها شناسایی شدند")
 
         segments = run_blocking_with_timeout(
@@ -492,9 +454,15 @@ async def health():
 
 
 @app.post("/upload/start")
-async def upload_start():
+async def upload_start(request: Request):
     cleanup_expired_jobs()
-    job_id = new_job()
+    body = {}
+    try:
+        body = await request.json()
+    except Exception:
+        pass
+    single_speaker_mode = bool(body.get("single_speaker_mode", False))
+    job_id = new_job(single_speaker_mode)
     return {"job_id": job_id}