Spaces:
Opera10
/
Running on Zero

File size: 40,877 Bytes
c3da9f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
# ==========================================================================
# FILE: voice_engine.py  — موتورِ «کلون کردن صدا» (Vevo)، کاملاً روی CPU
# --------------------------------------------------------------------------
# این فایل همان کدِ اسپیسِ صدای قبلی (Sada) است، بدونِ رابطِ Gradio:
#   • همان پچ‌ها (pkg_resources، LangSegment، Amphion و transformers)
#   • همان مدل: amphion/Vevo — تبدیلِ timbre با flow matching در ۳۲ گام
#   • همان برشِ هوشمند از روی سکوت (smart_split_audio_vevo) و همان ادغام
#   • همان تعدادِ هسته: همه‌ی هسته‌های واقعیِ کانتینر (cgroup)، بدونِ GPU
# هیچ تنظیمی از مدل عوض نشده؛ کیفیتِ خروجی دقیقاً مثلِ قبل است.
#
# چرا یک پروسه‌ی جدا (و نه داخلِ همان پروسه‌ی Gradio)؟
#   در این اسپیس «ساخت آهنگ» روی ZeroGPU است. ZeroGPU برای هر کارِ GPU از
#   پروسه‌ی اصلی fork می‌گیرد. اگر تبدیلِ صدا با ۱۶ رشته‌ی CPU داخلِ همان پروسه
#   اجرا شود، fork وسطِ کارِ چندرشته‌ایِ OpenMP/MKL می‌تواند پروسه‌ی GPU را قفل
#   کند. پچ‌های سراسریِ این موتور (LlamaConfig، pkg_resources و ...) هم نباید
#   روی مدل‌های آهنگ اثر بگذارند. برای همین این موتور در پروسه‌ی خودش اجرا
#   می‌شود و spaces/ZeroGPU را اصلاً ایمپورت نمی‌کند؛ یعنی دقیقاً مثلِ اسپیسِ
#   صدای قبلی، فقط CPU و RAM مصرف می‌کند و هیچ سهمیه‌ی GPU نمی‌گیرد.
#
# اجرا: voice_clone.py این فایل را با  python -u voice_engine.py --serve <socket>
# بالا می‌آورد و کارها را از راهِ یک سوکتِ محلی (فقط داخلِ همین کانتینر) می‌فرستد.
#
# 🆕 تغییرهای این نسخه (طبقِ درخواست و از روی لاگ‌های اسپیسِ قبلی):
#   ۱) صدای مرجع حداکثر ۱۲ ثانیه (VOICE_REF_MAX_SECONDS): مدلِ Vevo صدای مرجع را جلوی «هر» تکه
#      می‌گذارد؛ مرجعِ ۳ دقیقه‌ای یعنی هر تکه‌ی ۶ ثانیه‌ای مثلِ یک صدای ۳ دقیقه‌ای حساب می‌شد
#      (همان ۲۰۰ ثانیه‌ای‌های لاگ). مدل با مرجعِ حداکثر ~۱۲ ثانیه آموزش دیده؛ پس از مرجعِ بلندتر
#      «پُرصداترین» ۱۲ ثانیه برداشته می‌شود (از روی سکوت بریده، نه وسطِ کلمه). مرجعِ کوتاه‌تر
#      دست نمی‌خورد.
#   ۲) دو کارِ هم‌زمان (VOICE_CONCURRENT_JOBS=2، فقط وقتی اکشن پرو جای خالیِ دیگری ندارد):
#      کارِ تنها همه‌ی هسته‌ها را می‌گیرد (مثلِ قبل)؛ وقتی دو کار با هم‌اند هر کدام نصف. تقسیمِ
#      هسته‌ها در «هر گامِ» مدل دوباره حساب می‌شود (نه فقط اولِ کار).
#   ۳) کارِ بی‌صاحب لغو می‌شود: اگر درخواست‌دهنده دیگر منتظر نبود (اتصال قطع شد / زمانش تمام شد)،
#      همان لحظه (در گامِ بعدیِ مدل) متوقف می‌شود و CPU را برای کارهای واقعی آزاد می‌کند. قبلاً
#      کارِ بی‌صاحب تا آخر ادامه می‌داد و کارهای بعدی پشتش می‌ماندند (BrokenPipe و ۶۴۵/۶۴۷ ثانیه‌ای‌های
#      لاگ برای فایل‌های دوتکه‌ای).
# هر تکه دقیقاً مثلِ قبل ساخته می‌شود (همان نرمال‌سازی، همان ۳۲ گام، همان مدل) و تکه‌های یک
# فایل مثلِ قبل یکی‌یکی و به ترتیب ساخته و به هم وصل می‌شوند.
# ==========================================================================
import os
import sys
import json
import time
import signal
import threading
import importlib.util
import torch
import torchaudio  # noqa: F401  (مثلِ قبل: Amphion به آن نیاز دارد)
import numpy as np
from huggingface_hub import snapshot_download, hf_hub_download
import subprocess
import uuid
import soundfile as sf
import librosa
import shutil
import re
import warnings

# مرجعی که PySoundFile نمی‌خواند (بعضی mp3ها) خودکار با audioread خوانده می‌شود؛ فقط هشدارِ
# «منسوخ شدن» در لاگ می‌ماند که هیچ خطایی نیست — بی‌صدا شد تا لاگ تمیز بماند (رفتار همان است).
warnings.filterwarnings("ignore", message="PySoundFile failed")
warnings.filterwarnings("ignore", category=FutureWarning, module="librosa")

# --- رفع خطای pkg_resources (نسخه‌های جدید setuptools آن را حذف کرده‌اند و pyworld به آن نیاز دارد) ---
try:
    import pkg_resources  # noqa: F401
except ImportError:
    import types
    from importlib import metadata as _md

    class _DistributionNotFound(Exception):
        pass

    class _Dist:
        def __init__(self, name):
            try:
                self.version = _md.version(name)
            except _md.PackageNotFoundError:
                raise _DistributionNotFound(name)
            self.project_name = name
            self.key = name.lower()

    _pr = types.ModuleType("pkg_resources")
    _pr.DistributionNotFound = _DistributionNotFound
    _pr.get_distribution = lambda name: _Dist(name)
    _pr.require = lambda *a, **k: []
    sys.modules["pkg_resources"] = _pr
    print("[FIX] pkg_resources shim installed")

# --- 1. نصب و تنظیمات اولیه (بدون تغییر) ---

def install_espeak():
    try:
        result = subprocess.run(["which", "espeak-ng"], capture_output=True, text=True)
        if result.returncode != 0:
            print("Installing espeak-ng...")
            subprocess.run(["apt-get", "update"], check=True)
            subprocess.run(["apt-get", "install", "-y", "espeak-ng", "espeak-ng-data"], check=True)
    except Exception as e:
        print(f"Error installing espeak-ng: {e}")

install_espeak()

def patch_langsegment_init():
    try:
        spec = importlib.util.find_spec("LangSegment")
        if spec is None or spec.origin is None: return
        init_path = os.path.join(os.path.dirname(spec.origin), '__init__.py')
        with open(init_path, 'r') as f: lines = f.readlines()
        modified = False
        new_lines =[]
        target_line_prefix = "from .LangSegment import"
        for line in lines:
            if line.strip().startswith(target_line_prefix) and ('setLangfilters' in line or 'getLangfilters' in line):
                mod_line = line.replace(',setLangfilters', '').replace(',getLangfilters', '')
                mod_line = mod_line.replace('setLangfilters,', '').replace('getLangfilters,', '').rstrip(',')
                new_lines.append(mod_line + '\n')
                modified = True
            else:
                new_lines.append(line)
        if modified:
            with open(init_path, 'w') as f: f.writelines(new_lines)
            try:
                import LangSegment
                importlib.reload(LangSegment)
            except: pass
    except: pass

patch_langsegment_init()

if not os.path.exists("Amphion"):
    print("Cloning Amphion repository...")
    subprocess.run(["git", "clone", "https://github.com/open-mmlab/Amphion.git"])

amphion_path = os.path.abspath("Amphion")
if amphion_path not in sys.path:
    sys.path.append(amphion_path)

# --- پچ کردن تمام کدهای Amphion برای سازگاری با نسخه جدید transformers ---
def patch_amphion_llama_config():
    try:
        for root, dirs, files in os.walk(amphion_path):
            for file in files:
                if file.endswith(".py"):
                    file_path = os.path.join(root, file)
                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            content = f.read()

                        new_content = re.sub(
                            r'LlamaConfig\(\s*0\s*,\s*256\s*,\s*1024\s*,\s*1\s*,\s*1\s*\)',
                            r'LlamaConfig(vocab_size=0, hidden_size=256, intermediate_size=1024, num_hidden_layers=1, num_attention_heads=1)',
                            content
                        )

                        if new_content != content:
                            with open(file_path, "w", encoding="utf-8") as f:
                                f.write(new_content)
                            print(f"Patched {file} successfully.")
                    except Exception:
                        pass
    except Exception as e:
        print(f"Failed to patch Amphion: {e}")

patch_amphion_llama_config()

# --- رفع مشکل rope_theta با تزریق (Monkey Patch) در LlamaConfig ---
# (این پچ فقط داخلِ همین پروسه است و به مدل‌های ساخت آهنگ نمی‌رسد.)
import transformers
from transformers.models.llama.configuration_llama import LlamaConfig

_original_llama_init = LlamaConfig.__init__
def _patched_llama_init(self, *args, **kwargs):
    kwargs.setdefault('rope_theta', 10000.0)
    kwargs.setdefault('attention_bias', False)
    kwargs.setdefault('rope_scaling', None)
    kwargs.setdefault('max_position_embeddings', 4096)

    _original_llama_init(self, *args, **kwargs)

    if not hasattr(self, 'rope_theta'):
        self.rope_theta = 10000.0
    if not hasattr(self, 'attention_bias'):
        self.attention_bias = False

LlamaConfig.__init__ = _patched_llama_init

os.makedirs("wav", exist_ok=True)
os.makedirs("ckpts/Vevo", exist_ok=True)

try:
    from models.vc.vevo.vevo_utils import VevoInferencePipeline
except ImportError:
    sys.path.append(os.path.join(amphion_path))
    from models.vc.vevo.vevo_utils import VevoInferencePipeline

# --- توابع ذخیره و تنظیمات مدل ---

def save_audio_pcm16(waveform, output_path, sample_rate=24000):
    try:
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.detach().cpu()
            if waveform.dim() == 2 and waveform.shape[0] == 1:
                waveform = waveform.squeeze(0)
            waveform = waveform.numpy()
        sf.write(output_path, waveform, sample_rate, subtype='PCM_16')
    except Exception as e:
        print(f"Save error: {e}")

downloaded_resources = {
    "configs": False,
    "tokenizer_vq8192": False,
    "fmt_Vq8192ToMels": False,
    "vocoder": False
}

def setup_configs():
    if downloaded_resources["configs"]: return
    target_config_path = "models/vc/vevo/config"
    os.makedirs(target_config_path, exist_ok=True)

    config_files =["Vq8192ToMels.json", "Vocoder.json", "hubert_large_l18_mean_std.npz"]
    source_config_path = os.path.join(amphion_path, "models/vc/vevo/config")

    for file in config_files:
        target_file_path = f"{target_config_path}/{file}"
        source_file_path = os.path.join(source_config_path, file)

        if not os.path.exists(target_file_path):
            if os.path.exists(source_file_path):
                shutil.copy(source_file_path, target_file_path)
            else:
                try:
                    file_data = hf_hub_download(repo_id="amphion/Vevo", filename=f"config/{file}", repo_type="model")
                    if os.path.exists(target_file_path): os.remove(target_file_path)
                    subprocess.run(["cp", file_data, target_file_path])
                except Exception as e:
                    print(f"Error downloading config {file}: {e}")

    downloaded_resources["configs"] = True

setup_configs()

# =========================================================================
#  اجرای کامل روی CPU (بدون مصرف اعتبار ZeroGPU)
#  نکته: این پروسه spaces را ایمپورت نمی‌کند، پس GPUی در کار نیست؛ با این حال
#  دستگاه مثلِ قبل صریحاً CPU گذاشته می‌شود. کیفیت و تنظیمات مدل دست نخورده است.
# =========================================================================
def get_real_cpu_count():
    n = os.cpu_count() or 1
    try:
        n = min(n, len(os.sched_getaffinity(0)))
    except Exception:
        pass
    try:
        # سهمیه واقعی کانتینر (cgroup v2)
        with open("/sys/fs/cgroup/cpu.max") as f:
            quota, period = f.read().split()
            if quota != "max":
                n = min(n, max(1, int(int(quota) / int(period))))
    except Exception:
        pass
    return max(1, n)

CPU_THREADS = get_real_cpu_count()
# حداکثر چند درخواست هم‌زمان (۲ = «کارِ دوم هم‌زمان»؛ ۱ = رفتارِ قبلی). اکشن پرو کارِ دوم را فقط
# وقتی می‌فرستد که هیچ اسپیسِ صدای خالیِ دیگری نمانده باشد.
VOICE_CONCURRENT_JOBS = max(1, min(4, int(os.environ.get("VOICE_CONCURRENT_JOBS", "2") or 2)))
if CPU_THREADS < 2 * VOICE_CONCURRENT_JOBS:
    VOICE_CONCURRENT_JOBS = max(1, CPU_THREADS // 2)
torch.set_num_threads(CPU_THREADS)
try:
    torch.set_num_interop_threads(max(1, min(4, CPU_THREADS // 4)))
except Exception:
    pass
os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
print(f"[CPU MODE] Using {CPU_THREADS} CPU threads (os.cpu_count={os.cpu_count()}); "
      f"up to {VOICE_CONCURRENT_JOBS} job(s) at once (a lone job gets all threads)")

device = torch.device("cpu")
inference_pipelines = {}

downloaded_content_style_tokenizer_path = None
downloaded_fmt_path = None
downloaded_vocoder_path = None

def preload_all_resources():
    setup_configs()
    global downloaded_content_style_tokenizer_path, downloaded_fmt_path, downloaded_vocoder_path
    if not downloaded_resources["tokenizer_vq8192"]:
        downloaded_content_style_tokenizer_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["tokenizer/vq8192/"])
        downloaded_resources["tokenizer_vq8192"] = True
    if not downloaded_resources["fmt_Vq8192ToMels"]:
        downloaded_fmt_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["acoustic_modeling/Vq8192ToMels/"])
        downloaded_resources["fmt_Vq8192ToMels"] = True
    if not downloaded_resources["vocoder"]:
        downloaded_vocoder_path = snapshot_download(repo_id="amphion/Vevo", repo_type="model", cache_dir="./ckpts/Vevo", allow_patterns=["acoustic_modeling/Vocoder/*"])
        downloaded_resources["vocoder"] = True

preload_all_resources()

_pipeline_lock = threading.Lock()

def get_pipeline():
    with _pipeline_lock:
        if "timbre" in inference_pipelines: return inference_pipelines["timbre"]
        tokenizer_path = os.path.join(downloaded_content_style_tokenizer_path, "tokenizer/vq8192")
        fmt_ckpt = os.path.join(downloaded_fmt_path, "acoustic_modeling/Vq8192ToMels")
        vocoder_ckpt = os.path.join(downloaded_vocoder_path, "acoustic_modeling/Vocoder")
        pipeline = VevoInferencePipeline(
            content_style_tokenizer_ckpt_path=tokenizer_path,
            fmt_cfg_path="./models/vc/vevo/config/Vq8192ToMels.json",
            fmt_ckpt_path=fmt_ckpt,
            vocoder_cfg_path="./models/vc/vevo/config/Vocoder.json",
            vocoder_ckpt_path=vocoder_ckpt,
            device=device,
        )
        _install_step_hooks(pipeline)
        inference_pipelines["timbre"] = pipeline
        return pipeline


# =========================================================================
#  🆕 تقسیمِ هسته‌ها بینِ کارهای هم‌زمان + لغوِ کارِ بی‌صاحب، در «هر گامِ» مدل
# -------------------------------------------------------------------------
# روی ماژول‌های سنگینِ مدل (هر گامِ flow matching، لایه‌های HuBERT، vocoder ...) یک
# forward-pre-hook گذاشته می‌شود که در همان رشته‌ی کار اجرا می‌شود و دو کار می‌کند:
#   • اگر کار لغو شده (درخواست‌دهنده دیگر منتظر نیست) همان لحظه متوقفش می‌کند؛
#   • تعدادِ رشته‌های همین کار را با «همه‌ی هسته‌ها ÷ تعدادِ کارهای در جریان» یکی می‌کند
#     (torch.set_num_threads برای هر رشته جداست). کارِ تنها = همه‌ی هسته‌ها، مثلِ قبل.
# هیچ عددی از خودِ مدل عوض نمی‌شود؛ فقط این‌که هر لحظه چند هسته کار کنند.
# =========================================================================
_tls = threading.local()
_jobs_lock = threading.Lock()
_active_jobs = 0


class _JobCancelled(Exception):
    pass


class _Job:
    def __init__(self):
        self.cancelled = threading.Event()
        self.done = threading.Event()
        self.reason = ""
        self.threads = 0
        self.max_parallel = 1

    def cancel(self, reason):
        if self.done.is_set() or self.cancelled.is_set():
            return
        self.reason = reason
        self.cancelled.set()

    def check(self):
        if self.cancelled.is_set():
            raise _JobCancelled(self.reason or "cancelled")


def _sync_threads(job):
    """رشته‌های کارِ همین رشته را با تعدادِ کارهای در جریان هماهنگ می‌کند."""
    active = max(1, _active_jobs)
    if active > job.max_parallel:
        job.max_parallel = active
    want = max(1, CPU_THREADS // active)
    if job.threads != want:
        torch.set_num_threads(want)
        job.threads = want


def _step_hook(module, args):
    job = getattr(_tls, "job", None)
    if job is not None:
        job.check()
        _sync_threads(job)
    return None


def _install_step_hooks(pipeline):
    """هوک روی ماژول‌هایی که در هر تکه بارها صدا زده می‌شوند (و لایه‌های داخلِ ModuleListهایشان)."""
    names = ("fmt_model", "hubert_model", "vocoder_model", "mel_model", "content_style_tokenizer")
    hooked, roots = 0, []
    for name in names:
        root = getattr(pipeline, name, None)
        if not isinstance(root, torch.nn.Module):
            continue
        roots.append(name)
        targets = {id(root): root}
        estimator = getattr(root, "diff_estimator", None)
        if isinstance(estimator, torch.nn.Module):
            targets[id(estimator)] = estimator          # هر گامِ flow matching
        for sub in root.modules():
            if isinstance(sub, torch.nn.ModuleList):
                for layer in sub:
                    if isinstance(layer, torch.nn.Module):
                        targets[id(layer)] = layer      # لایه‌ها / بلوک‌ها
        for mod in targets.values():
            mod.register_forward_pre_hook(_step_hook)
            hooked += 1
    print(f"[CPU MODE] per-step thread/cancel hooks on {hooked} modules ({', '.join(roots) or 'none'})", flush=True)


# --- تابع جدید برش هوشمند با قابلیت افزایش داینامیک بازه جستجو ---
# ⚠️ اکشن پرو هم دقیقاً همین الگوریتم را (با همین عددها) پیش از ارسال اجرا
# می‌کند تا تکه‌ها را هم‌زمان به چند اسپیس بفرستد؛ اگر این‌جا عوض شد، آن‌جا
# (proxy_worker.py → _voice_split_points) هم باید عوض شود.
def smart_split_audio_vevo(y, sr, min_segment=4):
    total_samples = len(y)
    current_sample = 0
    chunks = []

    # اگر طول کل فایل کمتر از حد مجاز اولیه (8 ثانیه) باشد، نیازی به برش نیست
    if total_samples <= 8 * sr:
        return [y]

    while current_sample < total_samples:
        # اگر حجم نمونه‌های باقی‌مانده کمتر از حداقل زمان قطعه (4 ثانیه) باشد، بقیه فایل را یکجا بردار
        if total_samples - current_sample <= min_segment * sr:
            chunks.append(y[current_sample:])
            break

        # بازه‌های گام‌به‌گام برای جستجوی سکوت: ابتدا ۸، سپس ۱۲، ۱۵ و نهایتاً ۲۰ ثانیه
        steps = [8, 12, 15, 20]
        silence_threshold = 0.015  # آستانه شدت انرژی سیگنال برای پذیرش به عنوان سکوت واقعی

        best_split_point = None
        best_rms_val = float('inf')

        for max_sec in steps:
            start_search = current_sample + int(min_segment * sr)
            end_search = current_sample + int(max_sec * sr)

            if start_search >= total_samples:
                best_split_point = total_samples
                break

            if end_search > total_samples:
                end_search = total_samples

            search_region = y[start_search:end_search]
            if len(search_region) == 0:
                best_split_point = total_samples
                break

            # محاسبه میزان انرژی (RMS) قطعه
            rms = librosa.feature.rms(y=search_region, frame_length=2048, hop_length=512)[0]
            if len(rms) == 0:
                continue

            min_rms_idx = np.argmin(rms)
            min_rms_val = rms[min_rms_idx]
            candidate_point = start_search + (min_rms_idx * 512)

            # ذخیره بهترین نقطه کم‌صدا (در صورتی که سکوت مطلق پیدا نشود، کم‌صدا‌ترین نقطه انتخاب می‌شود)
            if min_rms_val < best_rms_val:
                best_rms_val = min_rms_val
                best_split_point = candidate_point

            # اگر شدت صدای نقطه پیدا شده پایین‌تر از آستانه سکوت بود، فرآیند جستجو را متوقف کرده و برش بزن
            if min_rms_val < silence_threshold:
                break

        if best_split_point is None:
            best_split_point = min(current_sample + int(8 * sr), total_samples)

        chunk_audio = y[current_sample:best_split_point]
        chunks.append(chunk_audio)
        current_sample = best_split_point

    # بررسی قطعه آخر: اگر بیش از حد کوتاه باشد، آن را با بخش قبلی ادغام می‌کنیم
    if len(chunks) > 1 and len(chunks[-1]) < int(1.5 * sr):
        last_chunk = chunks.pop()
        combined_len = len(chunks[-1]) + len(last_chunk)
        # تا سقف حداکثر ۲۰ ثانیه ادغام قطعه آخر مجاز است
        if combined_len <= int(20 * sr):
            chunks[-1] = np.concatenate([chunks[-1], last_chunk])
        else:
            chunks.append(last_chunk)

    return chunks


# =========================================================================
#  همان تابعِ اصلیِ قبلی (predict_voice_conversion)، فقط به‌جای gr.Progress
#  یک تابعِ گزارشِ پیشرفت می‌گیرد. خطاها مثلِ قبل بالا می‌روند و پروسه‌ی
#  Gradio همان پیامِ قبلی را («خطایی در حین پردازش رخ داد: ...») نشان می‌دهد.
#  تکه‌ها مثلِ قبل یکی‌یکی و به ترتیب ساخته و به هم وصل می‌شوند.
# =========================================================================
SR = 24000

# 🆕 سقفِ طولِ صدای مرجع (صدای مدل). توضیح: بالای فایل، بندِ ۱. ۰ = بدونِ سقف (رفتارِ قبلی).
REF_MAX_SECONDS = float(os.environ.get("VOICE_REF_MAX_SECONDS", "12"))
_REF_TOLERANCE_SECONDS = 0.5      # تا ۱۲٫۵ ثانیه دست نمی‌خورد (اکشن پرو خودش ۱۲ ثانیه می‌فرستد)
_REF_FRAME = 480                  # ۲۰ میلی‌ثانیه در ۲۴ کیلوهرتز (همان گامِ فریم‌های مدل)


def pick_reference_segment(y, sr=SR, max_seconds=REF_MAX_SECONDS):
    """صدای مرجعِ بلند ← «پُرصداترین» تکه‌ی حداکثر max_seconds ثانیه‌ای (بیشترین گفتار)، که درست
    پیش از شروعِ گفتار شروع و روی یک مکث تمام می‌شود (نه وسطِ کلمه). مرجعِ کوتاه‌تر از سقف بدونِ
    هیچ تغییری برمی‌گردد. (اکشن پرو همین را پیش از ارسال اجرا می‌کند: _voice_pick_reference.)"""
    y = np.asarray(y, dtype=np.float32)
    n = len(y)
    if max_seconds <= 0 or n <= int((max_seconds + _REF_TOLERANCE_SECONDS) * sr):
        return y
    frame = max(1, int(_REF_FRAME * sr / 24000))
    nf = n // frame
    win = max(1, int(max_seconds * sr / frame))
    if nf <= win:
        return y
    rms = np.sqrt(np.mean(np.square(y[: nf * frame].reshape(nf, frame), dtype=np.float64), axis=1))
    thr = max(0.008, 0.12 * float(np.percentile(rms, 95)))      # زیرِ این = مکث/سکوت
    active = (rms > thr).astype(np.int64)
    csum = np.concatenate([[0], np.cumsum(active)])
    step = max(1, int(0.5 * sr / frame))
    starts = np.arange(0, nf - win + 1, step)
    counts = csum[starts + win] - csum[starts]
    best = int(starts[int(np.argmax(counts))])                  # بیشترین گفتار؛ در تساوی، زودترین
    sec = max(1, int(sr / frame))
    # شروع: آخرین فریمِ ساکت در ±۱ ثانیه‌ی دورِ بهترین شروع (یعنی درست پیش از گفتار)
    h_lo, h_hi = max(0, best - sec), min(nf - win, best + sec) + 1
    s0 = h_lo + _last_quiet(rms[h_lo:h_hi], thr)
    # پایان: آخرین فریمِ ساکت در ۳ ثانیه‌ی آخرِ پنجره (تا جای ممکن بلند، روی مکث)
    t_lo, t_hi = max(s0 + 1, s0 + win - 3 * sec), min(nf, s0 + win)
    e0 = (t_lo + _last_quiet(rms[t_lo:t_hi], thr)) if t_hi > t_lo else t_hi
    e0 = max(e0, s0 + 1)
    return y[s0 * frame: e0 * frame]


def _last_quiet(values, thr):
    """اندیسِ آخرین فریمِ ساکت (زیرِ thr)؛ اگر نبود، آخرین فریمی که هم‌اندازه‌ی کم‌صداترین است."""
    values = np.asarray(values, dtype=np.float64)
    if len(values) == 0:
        return 0
    quiet = np.nonzero(values <= thr)[0]
    if len(quiet):
        return int(quiet[-1])
    low = float(values.min())
    return int(np.nonzero(values <= low * 1.05 + 1e-6)[0][-1])


def run_voice_conversion(source_audio_path, ref_audio_path, progress=None, job=None):
    if progress is None:
        progress = lambda value, desc="": None  # noqa: E731
    if job is None:
        job = _Job()

    task_id = str(uuid.uuid4())
    out_path = f"wav/{task_id}_out.wav"

    # پردازش صوت اصلی
    content_data, _ = librosa.load(source_audio_path, sr=SR)

    # پردازش صوت مرجع (🆕 مرجعِ بلندتر از سقف ← پُرصداترین تکه‌اش؛ کوتاه‌تر دست نمی‌خورد)
    ref_data, _ = librosa.load(ref_audio_path, sr=SR)
    ref_full = len(ref_data) / SR
    ref_data = pick_reference_segment(ref_data, SR)
    if len(ref_data) < int(ref_full * SR):
        print(f"[CPU MODE] reference {ref_full:.1f}s -> using its best {len(ref_data) / SR:.1f}s "
              f"(VOICE_REF_MAX_SECONDS={REF_MAX_SECONDS:g})", flush=True)
    ref_tensor = torch.FloatTensor(ref_data).unsqueeze(0)
    ref_tensor = ref_tensor / (torch.max(torch.abs(ref_tensor)) + 1e-6) * 0.95

    temp_r = f"wav/{task_id}_temp_r.wav"
    save_audio_pcm16(ref_tensor, temp_r, SR)

    try:
        # اعمال تقسیم‌کننده هوشمند چندمرحله‌ای
        chunks = smart_split_audio_vevo(content_data, SR, min_segment=4)

        generated_chunks = []
        pipeline = get_pipeline()

        # شبیه‌سازی قطعه به قطعه (مثلِ قبل، یکی‌یکی و به ترتیب)
        total_chunks = len(chunks)
        print(f"[CPU MODE] {total_chunks} chunks, {len(content_data)/SR:.1f}s audio "
              f"(ref {len(ref_data)/SR:.1f}s)", flush=True)
        for idx, chunk in enumerate(chunks):
            job.check()
            progress(idx / total_chunks, f"در حال پردازش روی CPU: قطعه {idx+1} از {total_chunks}")
            chunk_started = time.time()
            chunk_tensor = torch.FloatTensor(chunk).unsqueeze(0)
            chunk_tensor = chunk_tensor / (torch.max(torch.abs(chunk_tensor)) + 1e-6) * 0.95

            temp_c_chunk = f"wav/{task_id}_temp_c_{idx}.wav"
            save_audio_pcm16(chunk_tensor, temp_c_chunk, SR)

            try:
                _sync_threads(job)
                with torch.inference_mode():
                    gen = pipeline.inference_fm(
                        src_wav_path=temp_c_chunk,
                        timbre_ref_wav_path=temp_r,
                        flow_matching_steps=32,
                    )
            finally:
                if os.path.exists(temp_c_chunk):
                    os.remove(temp_c_chunk)

            if torch.isnan(gen).any():
                gen = torch.nan_to_num(gen, nan=0.0)
            gen_np = gen.detach().cpu().squeeze().numpy()

            generated_chunks.append(gen_np)
            if total_chunks > 1:
                print(f"   [CPU MODE] chunk {idx+1}/{total_chunks}: {len(chunk)/SR:.1f}s in "
                      f"{time.time() - chunk_started:.1f}s ({job.threads} threads)", flush=True)

        job.check()
        progress(1.0, "در حال ادغام و ذخیره خروجی")
        # ادغام قطعات خروجی (به ترتیب)
        final_gen = np.concatenate(generated_chunks, axis=0)

        # ذخیره نهایی
        save_audio_pcm16(final_gen, out_path, SR)
    finally:
        if os.path.exists(temp_r):
            os.remove(temp_r)

    return out_path


# =========================================================================
#  سرویس‌دهی به پروسه‌ی Gradio (فقط داخلِ همین کانتینر، با کلیدِ مخفی)
# =========================================================================
_job_slots = threading.BoundedSemaphore(VOICE_CONCURRENT_JOBS)
_state = {"busy": False, "warm": False, "jobs": 0, "cancelled": 0, "started_at": time.time()}


def _write_state(path, **extra):
    if not path:
        return
    data = {"pid": os.getpid(), "ready": True, "busy": _state["busy"], "warm": _state["warm"],
            "jobs": _state["jobs"], "cancelled": _state["cancelled"], "threads": CPU_THREADS,
            "concurrency": VOICE_CONCURRENT_JOBS, "active": _active_jobs, "ts": time.time()}
    data.update(extra)
    try:
        tmp = f"{path}.tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(data, f)
        os.replace(tmp, path)
    except Exception:
        pass


def _cleanup_old_outputs(max_age_seconds=2 * 3600):
    """خروجی‌های قدیمیِ پوشه‌ی wav پاک می‌شوند (Gradio هر خروجی را لحظه‌ی تحویل در
    کشِ خودش کپی می‌کند، پس بعد از چند ساعت این‌جا دیگر لازم نیستند)."""
    try:
        cutoff = time.time() - max_age_seconds
        for name in os.listdir("wav"):
            path = os.path.join("wav", name)
            try:
                if os.path.isfile(path) and os.path.getmtime(path) < cutoff:
                    os.remove(path)
            except Exception:
                pass
    except Exception:
        pass


def _watch_client(conn, job):
    """اگر درخواست‌دهنده (پروسه‌ی Gradio) اتصال را بست -- یعنی دیگر کسی منتظرِ این نتیجه نیست
    (کاربر/اکشن پرو رفت یا زمانش تمام شد) -- کار همان لحظه لغو می‌شود."""
    while not job.done.is_set() and not job.cancelled.is_set():
        try:
            ready = conn.poll(0.5)
        except (OSError, EOFError, ValueError):
            job.cancel("the requester closed the connection")
            return
        if not ready:
            continue
        try:
            msg = conn.recv()
        except (OSError, EOFError, ValueError):
            job.cancel("the requester is no longer waiting")
            return
        if (msg or {}).get("op") == "cancel":
            job.cancel("cancel requested")
            return


def _serve(sock_path, state_path):
    from multiprocessing.connection import Listener

    authkey = bytes.fromhex(os.environ.get("VOICE_ENGINE_AUTHKEY", ""))
    if not authkey:
        raise SystemExit("VOICE_ENGINE_AUTHKEY is missing")
    try:
        os.remove(sock_path)
    except FileNotFoundError:
        pass
    listener = Listener(sock_path, family="AF_UNIX", authkey=authkey)
    _write_state(state_path)
    print(f"🎙️ [Voice Engine] آماده است (pid={os.getpid()}, CPU threads={CPU_THREADS}, "
          f"تا {VOICE_CONCURRENT_JOBS} کارِ هم‌زمان).", flush=True)

    def handle(conn):
        global _active_jobs
        try:
            msg = conn.recv()
            op = (msg or {}).get("op")
            if op == "ping":
                conn.send({"op": "pong", "busy": _state["busy"], "warm": _state["warm"], "jobs": _state["jobs"],
                           "active": _active_jobs, "concurrency": VOICE_CONCURRENT_JOBS})
                return
            if op != "convert":
                conn.send({"op": "error", "message": f"unknown op: {op}"})
                return
            job = _Job()
            with _job_slots:
                with _jobs_lock:
                    _active_jobs += 1
                    _state["busy"] = True
                _write_state(state_path)
                threading.Thread(target=_watch_client, args=(conn, job), daemon=True,
                                 name="voice-client-watch").start()
                _tls.job = job
                started = time.time()
                send_lock = threading.Lock()

                def report(value, desc=""):
                    if job.cancelled.is_set():
                        return
                    try:
                        with send_lock:
                            conn.send({"op": "progress", "value": float(value), "desc": str(desc)})
                    except Exception:
                        job.cancel("the requester is no longer waiting")

                try:
                    out_path = run_voice_conversion(msg.get("source"), msg.get("ref"), progress=report, job=job)
                    _state["warm"] = True
                    job.check()
                    try:
                        with send_lock:
                            conn.send({"op": "done", "path": os.path.abspath(out_path)})
                    except Exception:
                        job.cancel("the requester left right before the result")
                        raise _JobCancelled(job.reason)
                    job.done.set()
                    together = " (هم‌زمان با کارِ دیگر)" if job.max_parallel > 1 else ""
                    print(f"✅ [Voice Engine] تبدیل در {time.time() - started:.1f} ثانیه کامل شد{together}.", flush=True)
                except _JobCancelled:
                    _state["cancelled"] += 1
                    print(f"⏹️ [Voice Engine] کار بعد از {time.time() - started:.1f} ثانیه لغو شد "
                          f"({job.reason}); CPU آزاد شد.", flush=True)
                except Exception as e:
                    if job.cancelled.is_set():
                        _state["cancelled"] += 1
                        print(f"⏹️ [Voice Engine] کار بعد از {time.time() - started:.1f} ثانیه لغو شد "
                              f"({job.reason}); CPU آزاد شد.", flush=True)
                    else:
                        print(f"❌ [Voice Engine] خطا: {type(e).__name__}: {e}", flush=True)
                        try:
                            with send_lock:
                                conn.send({"op": "error", "message": str(e)})
                        except Exception:
                            pass
                finally:
                    job.done.set()
                    _tls.job = None
                    with _jobs_lock:
                        _active_jobs = max(0, _active_jobs - 1)
                        _state["busy"] = _active_jobs > 0
                        _state["jobs"] += 1
                        jobs_done = _state["jobs"]
                    _write_state(state_path)
                    if jobs_done % 25 == 0:
                        _cleanup_old_outputs()
        except EOFError:
            pass
        except Exception as e:
            print(f"⚠️ [Voice Engine] اتصال: {type(e).__name__}: {e}", flush=True)
        finally:
            try:
                conn.close()
            except Exception:
                pass

    while True:
        try:
            conn = listener.accept()
        except Exception as e:
            print(f"⚠️ [Voice Engine] accept: {type(e).__name__}: {e}", flush=True)
            time.sleep(0.5)
            continue
        threading.Thread(target=handle, args=(conn,), daemon=True).start()


def _exit_with_parent(parent_pid):
    """اگر پروسه‌ی اصلیِ اسپیس بسته شد، این موتور هم بسته شود (پروسه‌ی یتیم نماند)."""
    while True:
        time.sleep(5)
        if os.getppid() != parent_pid:
            print("🛑 [Voice Engine] پروسه‌ی اصلی بسته شد؛ موتور صدا هم بسته می‌شود.", flush=True)
            os._exit(0)


def main():
    sock_path = sys.argv[sys.argv.index("--serve") + 1]
    state_path = sys.argv[sys.argv.index("--state") + 1] if "--state" in sys.argv else ""
    threading.Thread(target=_exit_with_parent, args=(os.getppid(),), daemon=True).start()
    signal.signal(signal.SIGTERM, lambda *_: os._exit(0))

    # گرم‌کردنِ مدل همان اول (قبلاً اولین درخواست این زمان را می‌داد). اگر نشد،
    # مثلِ قبل در اولین درخواست دوباره ساخته می‌شود.
    try:
        t0 = time.time()
        get_pipeline()
        _state["warm"] = True
        print(f"🔥 [Voice Engine] مدل Vevo در {time.time() - t0:.1f} ثانیه روی CPU آماده شد.", flush=True)
        _warm_up()
    except Exception as e:
        print(f"⚠️ [Voice Engine] گرم‌کردنِ مدل ناموفق بود ({type(e).__name__}: {e})؛ "
              f"در اولین درخواست دوباره امتحان می‌شود.", flush=True)
    _serve(sock_path, state_path)


def _warm_up():
    """🆕 یک تبدیلِ آزمایشیِ خیلی کوتاه پیش از اولین درخواست (صدای ساختگیِ ۲ ثانیه‌ای): اولین
    اجرای مدل و librosa کارهای یک‌باره‌ای دارد (کامپایل/آماده‌سازی) که قبلاً به اولین کاربر می‌افتاد
    (همان ۴۰ ثانیه‌ایِ اولِ لاگ برای یک تکه‌ی ۶ ثانیه‌ای). خروجی بلافاصله پاک می‌شود."""
    t0 = time.time()
    tag = uuid.uuid4().hex
    src, ref, out = f"wav/{tag}_warm_src.wav", f"wav/{tag}_warm_ref.wav", None
    try:
        t = np.arange(int(2.0 * SR)) / SR
        tone = (0.2 * np.sin(2 * np.pi * 180 * t) * np.hanning(len(t))).astype(np.float32)
        sf.write(src, tone, SR, subtype="PCM_16")
        sf.write(ref, tone[::-1].copy(), SR, subtype="PCM_16")
        out = run_voice_conversion(src, ref)
        print(f"🔥 [Voice Engine] گرم‌کردنِ اولیه در {time.time() - t0:.1f} ثانیه انجام شد.", flush=True)
    except Exception as e:
        print(f"⚠️ [Voice Engine] گرم‌کردنِ اولیه انجام نشد ({type(e).__name__}: {e}); مشکلی نیست.", flush=True)
    finally:
        for path in (src, ref, out):
            try:
                if path and os.path.exists(path):
                    os.remove(path)
            except Exception:
                pass


if __name__ == "__main__":
    main()