Viewing File: /home/ubuntu/codegamaai-test/voice_clone/src/save.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, HttpUrl
from typing import List
import requests
import os
import boto3
current_directory = os.getcwd()
user_config_files = f"{current_directory}/src/user_config_files/"
dataset_store = f"{current_directory}/data/"
model_weights_dir = f"{current_directory}/src/rvc_implementation/assets/weights/"
model_index_dir = f"{current_directory}/src/rvc_implementation/logs/"
model_base_dir = f"{current_directory}/data/"
new_directory = f"{current_directory}/src/rvc_implementation"


os.environ['user_config_files'] = user_config_files
os.environ['dataset_store'] = dataset_store
os.environ['model_weights_dir'] = model_weights_dir
os.environ['model_index_dir'] = model_index_dir
os.environ['model_base_dir'] = model_base_dir
os.environ['new_directory'] = new_directory

app = FastAPI()


class AudioPayload(BaseModel):
    user_id: str
    vid: str
    audio_files: List[HttpUrl]

# Set up AWS S3 access
s3 = boto3.client('s3')

def download_file_from_s3(s3_url, local_path):
    bucket = s3_url.split('/')[2]
    key = '/'.join(s3_url.split('/')[3:])
    s3.download_file(Bucket=bucket, Key=key, Filename=local_path)

@app.post("/upload_audio/")
async def download_audios(payload: AudioPayload):
    directory = os.path.join(os.environ['dataset_store'], payload.user_id, payload.vid, "audio")
    if not os.path.exists(directory):
        os.makedirs(directory)

    for url in payload.audio_files:
        try:
            filename = url.split('/')[-1]
            local_file_path = os.path.join(directory, filename)
            download_file_from_s3(url, local_file_path)
            print(f"Downloaded and saved {filename} to {directory}")
        except Exception as e:
            print(f"Failed to download {url}: {str(e)}")
            raise HTTPException(status_code=400, detail=f"Failed to download {url}: {str(e)}")
    
    return {"message": "Files downloaded successfully", "directory": directory}

# To run the server, use:
# uvicorn script_name:app --reload

# Run the server with:
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="localhost", port=8000)