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

# 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] Modelo no encontrado en S3, descargando desde Hugging Face a S3...")
                self.download_and_upload_from_huggingface(model_name)
                model_files = self.get_model_file_parts(model_prefix)
            
            # 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")
            
            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...")
            model_url = f"https://huggingface.co/{model_name}/resolve/main/"
            files_to_download = [
                f"{model_name}/pytorch_model.bin",
                f"{model_name}/config.json",
                f"{model_name}/tokenizer.json",
            ]
            
            for file in files_to_download:
                file_url = model_url + file
                s3_key = file
                self.download_and_upload_to_s3_url(file_url, 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 download_and_upload_to_s3_url(self, url: str, s3_key: str):
        try:
            print(f"[INFO] Descargando archivo desde {url}...")
            response = requests.get(url)
            if response.status_code == 200:
                print(f"[INFO] Subiendo archivo a S3 con key {s3_key}...")
                self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=response.content)
            else:
                print(f"[ERROR] Error al descargar el archivo desde {url}, código de estado {response.status_code}.")
                raise HTTPException(status_code=500, detail=f"Error al descargar el archivo desde {url}")
        except Exception as e:
            print(f"[ERROR] Error al procesar la URL {url}: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Error al procesar la URL {url}")

@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 {"input_text": input_text, "output": outputs}

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

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