app.py

import os
import boto3
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import safetensors.torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from fastapi.responses import StreamingResponse
import io
import requests
import uvicorn
from dotenv import load_dotenv
import re
from tqdm import tqdm
from huggingface_hub import hf_hub_download

# Cargar las variables de entorno desde el archivo .env
load_dotenv()

# Configuración de AWS y Hugging Face
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_REGION = os.getenv("AWS_REGION")
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")  # Nombre del bucket S3
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Token de Hugging Face

# Cliente S3 de Amazon
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_REGION
)

app = FastAPI()

# Pydantic Model para el cuerpo de la solicitud del endpoint /predict/
class DownloadModelRequest(BaseModel):
    model_name: str
    pipeline_task: str
    input_text: str

class S3DirectStream:
    def __init__(self, bucket_name):
        self.s3_client = boto3.client(
            's3',
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
            region_name=AWS_REGION
        )
        self.bucket_name = bucket_name

    def stream_from_s3(self, key):
        try:
            print(f"[INFO] Descargando archivo {key} desde S3...")
            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
            return response['Body']  # Devolver el cuerpo directamente para el StreamingResponse
        except self.s3_client.exceptions.NoSuchKey:
            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
        except Exception as e:
            print(f"[ERROR] Error al descargar {key}: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Error al descargar archivo {key} desde S3.")

    def file_exists_in_s3(self, key):
        try:
            self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
            return True
        except self.s3_client.exceptions.ClientError:
            return False

    def load_model_from_stream(self, model_name):
        try:
            print(f"[INFO] Cargando el modelo {model_name} desde S3...")
            model_prefix = model_name.lower()
            model_files = self.get_model_file_parts(model_prefix)

            if not model_files:
                print(f"[INFO] El modelo {model_name} no está en S3. Procediendo a descargar desde Hugging Face...")
                self.download_and_upload_from_huggingface(model_name)
                model_files = self.get_model_file_parts(model_prefix)

            if not model_files:
                raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")

            # Leer y cargar todos los archivos del modelo
            model_streams = []
            for model_file in tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
                model_streams.append(self.stream_from_s3(model_file))

            config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
            config_data = config_stream.read().decode("utf-8")

            # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
            if model_files[0].endswith("model.safetensors"):
                print("[INFO] Cargando el modelo como safetensor...")
                model = AutoModelForCausalLM.from_config(config_data)
                model.load_state_dict(safetensors.torch.load_stream(model_streams[0]))
            else:
                print("[INFO] Cargando el modelo como archivo binario de PyTorch...")
                model = AutoModelForCausalLM.from_config(config_data)
                model.load_state_dict(torch.load(model_streams[0], map_location="cpu"))

            print("[INFO] Modelo cargado con éxito.")
            return model

        except Exception as e:
            print(f"[ERROR] Error al cargar el modelo desde S3: {e}")
            raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.")

    def load_tokenizer_from_stream(self, model_name):
        try:
            print(f"[INFO] Cargando el tokenizer {model_name} desde S3...")
            tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
            return tokenizer
        except Exception as e:
            print(f"[ERROR] Error al cargar el tokenizer desde S3: {e}")
            raise HTTPException(status_code=500, detail="Error al cargar el tokenizer desde S3.")

    def get_model_file_parts(self, model_name):
        print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_name}...")
        files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
        model_files = []
        for obj in tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
            key = obj['Key']
            if re.match(rf"{model_name}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"):
                model_files.append(key)
        if not model_files:
            print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.")
        return model_files

    def download_and_upload_from_huggingface(self, model_name):
        try:
            print(f"[INFO] Descargando {model_name} desde Hugging Face...")
            # Descargar los archivos necesarios de Hugging Face usando huggingface_hub
            files_to_download = [
                "pytorch_model.bin",
                "config.json",
                "tokenizer.json",
            ]

            for file in files_to_download:
                # Descargar cada archivo desde Hugging Face y subirlo a S3
                file_path = hf_hub_download(repo_id=model_name, filename=file, use_auth_token=HUGGINGFACE_TOKEN)
                s3_key = f"{model_name}/{file}"
                self.upload_file_to_s3(file_path, s3_key)
        except Exception as e:
            print(f"[ERROR] Error al descargar y subir modelo desde Hugging Face: {e}")
            raise HTTPException(status_code=500, detail="Error al descargar y subir modelo desde Hugging Face.")

    def upload_file_to_s3(self, file_path, s3_key):
        try:
            print(f"[INFO] Subiendo archivo {file_path} a S3 con key {s3_key}...")
            with open(file_path, 'rb') as data:
                self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data)
        except Exception as e:
            print(f"[ERROR] Error al subir archivo a S3: {e}")
            raise HTTPException(status_code=500, detail="Error al subir archivo a S3.")

@app.post("/predict/")
async def predict(model_request: DownloadModelRequest):
    try:
        print(f"[INFO] Recibiendo solicitud para predecir con el modelo {model_request.model_name}...")
        streamer = S3DirectStream(S3_BUCKET_NAME)
        model = streamer.load_model_from_stream(model_request.model_name)
        tokenizer = streamer.load_tokenizer_from_stream(model_request.model_name)

        task = model_request.pipeline_task
        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-image", "text-to-audio", "text-to-video"]:
            raise HTTPException(status_code=400, detail="Pipeline task no soportado")

        nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, max_length=2046)

        input_text = model_request.input_text
        print(f"[INFO] Ejecutando tarea {task} con el texto de entrada...")
        outputs = nlp_pipeline(input_text)

        # Eliminación de archivo local después de subir a S3
        if task == "text-to-speech":
            s3_key = f"{model_request.model_name}/generated_audio.wav"
            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav")
        elif task == "text-to-image":
            s3_key = f"{model_request.model_name}/generated_image.png"
            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
        elif task == "text-to-video":
            s3_key = f"{model_request.model_name}/generated_video.mp4"
            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")

        return {"output": outputs}
    except Exception as e:
        print(f"[ERROR] Error al procesar la solicitud: {str(e)}")
        raise HTTPException(status_code=500, detail="Error al realizar la predicción.")

if __name__ == "__main__":
    print("Iniciando servidor de predicciones en localhost:8000")
    uvicorn.run(app, host="0.0.0.0", port=8000)