File size: 2,138 Bytes
9a3c8f5
 
932db78
 
9a3c8f5
 
932db78
 
 
9a3c8f5
932db78
 
9a3c8f5
 
d8c037d
 
 
 
 
fb74953
d8c037d
932db78
0309473
932db78
 
 
9a3c8f5
 
59da368
 
 
932db78
 
d8c037d
932db78
 
 
 
 
 
d8c037d
 
 
 
 
 
932db78
d8c037d
9a3c8f5
 
 
932db78
 
9a3c8f5
 
932db78
 
9a3c8f5
 
932db78
59da368
 
 
 
9a3c8f5
59da368
 
 
 
 
932db78
 
9a3c8f5
0309473
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
import io

from pymilvus import Collection
import asyncio
from bark import SAMPLE_RATE, generate_audio, preload_models
from scipy.io.wavfile import write
from db_connect import connect

# for audio storage
from audio_db.is3.is3 import StagedObject

connect()
preload_models()

# get response and audio response
response = Collection("Response")
audio_response = Collection("AudioResponse")

def check_if_exist(text):
    return len(audio_response.query(expr=f"text == \"{text}\"")) != 0


async def make_audio():

    data = []

    response_iterator = response.query_iterator(batch_size=1, output_fields=['text', 'embeddings'])
    ids_to_delete = []
    current_iteration = 0
    batch_size = 5

    while True:
        res = response_iterator.next()

        if len(res) == 0:
            print("query iteration finished, close")
            # close the iterator
            response_iterator.close()
            break

        text = res[0]['text']
        # remove the response to generate the audio for if it already exists in the audio database
        if check_if_exist(text):
            response.delete(expr=f"id in {str([res[0]['id']])}")
            continue

        # generate audio
        audio_array = generate_audio(text, history_prompt="en_speaker_3")
        bytes_io = io.BytesIO()
        write(bytes_io, SAMPLE_RATE, audio_array)
        audio_bytes = bytes_io.read()

        # store the audio
        obj = StagedObject(obj=audio_bytes, name='audio')
        uploaded_object = await obj.upload()

        # save the audio record to AudioResponse
        data.append([res[0]['text'], uploaded_object.obj_id, res[0]['embeddings']])
        ids_to_delete.append(res[0]['id'])

        current_iteration += 1
        if current_iteration % batch_size == 0:
            # insert the audio
            audio_response.insert(list(zip(*data)))

            # delete text to generate audio
            response.delete(expr=f"id in {str(ids_to_delete)}")
            ids_to_delete = []

    response_iterator.close()

if __name__ == '__main__':
    loop = asyncio.new_event_loop()
    loop.run_until_complete(make_audio())