Viewing File: /home/ubuntu/combine_ai/bg_upscale/src/download.py

import boto3
import os
from botocore.exceptions import ClientError
from dotenv import load_dotenv
from src.utils import *
import requests
import uuid


load_dotenv()
s3_client = boto3.client(
    's3',
    aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
    aws_secret_access_key= os.environ['AWS_SECRET_ACCESS_KEY'],
    region_name= os.environ['AWS_DEFAULT_REGION']
)

def generate_unique_filename():
    unique_id = str(uuid.uuid4())
    return unique_id


def download_file_from_s3(s3_url, local_dir):
    parts = s3_url.split('/')
    bucket_name = parts[2].split('.')[0]
    s3_key = '/'.join(parts[3:])
    file_name = s3_key.split('/')[-1]

    local_file_path = os.path.join(local_dir, file_name)

    try:
        s3_client.download_file(bucket_name, s3_key, local_file_path)
        print(f"File downloaded successfully: {local_file_path}")
        return local_file_path
    except ClientError as e:
        print(f"Failed to download file from S3: {e}")
        return None
    
def direct_download(url, local_dir):
    local_file_path = os.path.join(local_dir, url.split('/')[-1])
    print("local_file_path", local_file_path)
    try:
        response = requests.get(url)
        with open(local_file_path, 'wb') as f:
            f.write(response.content)
        print(f"File downloaded successfully: {local_file_path}")
        return local_file_path
    except Exception as e:
        print(f"Failed to download file: {e}")
        return None
    

def direct_save(url, user_id, vid):
    # Create the directory if it does not exist
    directory = os.path.join(os.environ['model_base_dir'], user_id, vid, "raw_audio")
    if not os.path.exists(directory):
        os.makedirs(directory)
    local_file_path = os.path.join(os.environ['model_base_dir'], user_id, vid, "raw_audio", url.split('/')[-1])

    try:
        response = requests.get(url)
        with open(local_file_path, 'wb') as f:
            f.write(response.content)
        print(f"File downloaded successfully: {local_file_path}")
        return local_file_path
    except Exception as e:
        print(f"Failed to download file: {e}")
        return None
    

    

def image_output_file(user_id, image_output, type):
    
    image_name = generate_unique_filename() + ".jpg"
    # Save the image to the directory

    image_output_file_path = os.path.join(os.environ['REMOVE_BACKGROUND_DIRECTORY'], 'data', str(type), 'app_data', str(user_id), image_output)

    live_path = os.path.join(os.environ['OUTPUT_MEDIA_URL'], 'data', str(type), 'app_data', str(user_id), image_name)
    return image_output_file_path, live_path




img_path = "https://file-examples.com/storage/fed5266c9966708dcaeaea6/2017/10/file_example_JPG_100kB.jpg"
user_id = "remiko"
input_dir = "/home/zrlhowsqpnco/codegama_bot/combine_ai/bg_upscale/removebg/data/user_data/input_data/{user_id}"

# print(direct_download(img_path, input_dir)) # src/download.py