Hugging Face's logo

Spaces:
cocktailpeanut
/
maskgct
Running

App Files Community
1
maskgct
File size: 4,753 Bytes 
c968fc3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


# This code is modified from
# https://github.com/lifeiteng/vall-e/blob/9c69096d603ce13174fb5cb025f185e2e9b36ac7/valle/data/input_strategies.py
import random
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from typing import Tuple, Type

from lhotse import CutSet
from lhotse.dataset.collation import collate_features
from lhotse.dataset.input_strategies import (
    ExecutorType,
    PrecomputedFeatures,
    _get_executor,
)
from lhotse.utils import fastcopy


class PromptedFeatures:
    def __init__(self, prompts, features):
        self.prompts = prompts
        self.features = features

    def to(self, device):
        return PromptedFeatures(self.prompts.to(device), self.features.to(device))

    def sum(self):
        return self.features.sum()

    @property
    def ndim(self):
        return self.features.ndim

    @property
    def data(self):
        return (self.prompts, self.features)


class PromptedPrecomputedFeatures(PrecomputedFeatures):
    def __init__(
        self,
        dataset: str,
        cuts: CutSet,
        num_workers: int = 0,
        executor_type: Type[ExecutorType] = ThreadPoolExecutor,
    ) -> None:
        super().__init__(num_workers, executor_type)
        self.utt2neighbors = self._create_utt2neighbors(dataset, cuts)

    def __call__(self, cuts: CutSet) -> Tuple[PromptedFeatures, PromptedFeatures]:
        features, features_lens = self._collate_features(cuts)
        prompts, prompts_lens = self._collate_prompts(cuts)
        return PromptedFeatures(prompts, features), PromptedFeatures(
            prompts_lens, features_lens
        )

    def _create_utt2neighbors(self, dataset, cuts):
        utt2neighbors = defaultdict(lambda: [])
        utt2cut = {cut.id: cut for cut in cuts}
        if dataset.lower() == "libritts":
            self._process_libritts_dataset(utt2neighbors, utt2cut, cuts)
        elif dataset.lower() == "ljspeech":
            self._process_ljspeech_dataset(utt2neighbors, utt2cut, cuts)
        else:
            raise ValueError("Unsupported dataset")
        return utt2neighbors

    def _process_libritts_dataset(self, utt2neighbors, utt2cut, cuts):
        speaker2utts = defaultdict(lambda: [])
        for cut in cuts:
            speaker = cut.supervisions[0].speaker
            speaker2utts[speaker].append(cut.id)

        for spk, uttids in speaker2utts.items():
            sorted_uttids = sorted(uttids)
            if len(sorted_uttids) == 1:
                utt2neighbors[sorted_uttids[0]].append(utt2cut[sorted_uttids[0]])
                continue

            utt2prevutt = dict(
                zip(sorted_uttids, [sorted_uttids[1]] + sorted_uttids[:-1])
            )
            utt2postutt = dict(zip(sorted_uttids[:-1], sorted_uttids[1:]))
            for utt in sorted_uttids:
                if utt in utt2prevutt:
                    utt2neighbors[utt].append(utt2cut[utt2prevutt[utt]])
                if utt in utt2postutt:
                    utt2neighbors[utt].append(utt2cut[utt2postutt[utt]])

    def _process_ljspeech_dataset(self, utt2neighbors, utt2cut, cuts):
        uttids = [cut.id for cut in cuts]
        if len(uttids) == 1:
            utt2neighbors[uttids[0]].append(utt2cut[uttids[0]])
            return

        utt2prevutt = dict(zip(uttids, [uttids[1]] + uttids[:-1]))
        utt2postutt = dict(zip(uttids[:-1], uttids[1:]))
        for utt in uttids:
            prevutt, postutt = utt2prevutt.get(utt), utt2postutt.get(utt)
            if prevutt and utt[:5] == prevutt[:5]:
                utt2neighbors[utt].append(utt2cut[prevutt])
            if postutt and utt[:5] == postutt[:5]:
                utt2neighbors[utt].append(utt2cut[postutt])

    def _collate_features(self, cuts):
        return collate_features(
            cuts,
            executor=_get_executor(self.num_workers, executor_type=self._executor_type),
        )

    def _collate_prompts(self, cuts):
        prompts_cuts = []
        for k, cut in enumerate(cuts):
            prompts_cut = random.choice(self.utt2neighbors[cut.id])
            prompts_cuts.append(fastcopy(prompts_cut, id=f"{cut.id}-{str(k)}"))

        mini_duration = min([cut.duration for cut in prompts_cuts] + [3.0])
        prompts_cuts = CutSet(
            cuts={k: cut for k, cut in enumerate(prompts_cuts)}
        ).truncate(max_duration=mini_duration, offset_type="random", preserve_id=False)

        return collate_features(
            prompts_cuts,
            executor=_get_executor(self.num_workers, executor_type=self._executor_type),
        )