hugvey/hugvey/voice.py

import os
import time
import json
import logging
import threading
#from requests_threads import AsyncSession
from hashlib import sha1
import asyncio
from tornado.httpclient import AsyncHTTPClient, HTTPRequest

mainLogger = logging.getLogger("hugvey")
logger = mainLogger.getChild("voice")

class VoiceStorage(object):
    """
    Store & keep voices that are not part of the story json
    """
    def __init__(self, cache_dir, token):
        self.cache_dir = cache_dir
        if not os.path.exists(self.cache_dir):
            raise Exception(f"Cache dir does not exists: {self.cache_dir}")
#         self.request_session = AsyncSession(n=5)
        self.pendingRequests = {}
        self.token = token
    
    def getId(self, text):
        return sha1(text.encode()).hexdigest()
    
    def getFilename(self, text, isVariable=False):
        subdir = 'static' if not isVariable else 'variable'
        id = self.getId(text)
        prefix = id[:2]
        storageDir = os.path.join(self.cache_dir, subdir, prefix)
        fn = os.path.join(storageDir, f"{id}.wav")
        return fn
    
    async def requestFile(self, text, isVariable=False) -> str:
        id = self.getId(text)
        fn = self.getFilename(text)
        
        if os.path.exists(fn):
            return fn
        
        if id in self.pendingRequests:
            await self.pendingRequests[id].wait()
            if os.path.exists(fn):
                return fn
            return None
        
        dirname = os.path.dirname(fn)
        if not os.path.exists(dirname):
            logger.debug(f"create directory for file: {dirname}")
            os.makedirs(dirname, exist_ok=True)
            
        self.pendingRequests[id] = asyncio.Event()
        
        http_client = AsyncHTTPClient()
        request = HTTPRequest(
            method="POST",
            url="https://avatar.lyrebird.ai/api/v0/generate",
            body=json.dumps({"text": text}),
            headers={"authorization": f"Bearer {self.token}"}
        )
        try:
            response = await http_client.fetch(request)
        except Exception as e:
            logger.exception(e)
            self.pendingRequests[id].set()
            return None
        else:
            if response.code != 200:
                logger.critical(f"No proper response! {response.code}")
                self.pendingRequests[id].set()
                return None
            
            logger.debug(f"Wrote body: {response.code}")
            with open(fn, "wb") as f:
                f.write(response.body)
            self.pendingRequests[id].set()
            print(type(fn), fn)
            return fn