app.py

import os
import logging
import boto3
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import hf_hub_download
import asyncio

# Configuración de variables
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")
HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")

MAX_TOKENS = 1024

# Configuración de cliente S3
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
)

# Inicialización de la app FastAPI
app = FastAPI()

# Estructura de solicitudes
class GenerateRequest(BaseModel):
    model_name: str
    input_text: str
    task_type: str

# Clase para manejo de S3
class S3Manager:
    def __init__(self, bucket_name):
        self.bucket_name = bucket_name
        self.s3_client = s3_client

    async def get_file(self, key: str):
        """Descarga un archivo desde S3."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._get_file_sync, key)

    def _get_file_sync(self, key: str):
        try:
            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
            return response['Body'].read()
        except self.s3_client.exceptions.NoSuchKey:
            raise HTTPException(status_code=404, detail=f"Archivo {key} no encontrado en S3.")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Error al obtener el archivo {key} de S3: {str(e)}")

    async def upload_file(self, file_path: str, key: str):
        """Sube un archivo a S3."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._upload_file_sync, file_path, key)

    def _upload_file_sync(self, file_path: str, key: str):
        try:
            with open(file_path, "rb") as file:
                self.s3_client.put_object(Bucket=self.bucket_name, Key=key, Body=file)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Error al subir {key} a S3: {str(e)}")

    async def file_exists(self, key: str):
        """Verifica si un archivo existe en S3."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._file_exists_sync, key)

    def _file_exists_sync(self, key: str):
        try:
            self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
            return True
        except self.s3_client.exceptions.ClientError:
            return False
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Error al verificar existencia de {key}: {str(e)}")

    async def download_model_files(self, model_name: str):
        """Descarga los archivos del modelo desde Hugging Face y los sube a S3 si no están presentes."""
        model_name_s3 = model_name.replace("/", "-").lower()
        files = ["pytorch_model.bin", "tokenizer.json", "config.json"]

        for file in files:
            if not await self.file_exists(f"{model_name_s3}/{file}"):
                local_file = hf_hub_download(repo_id=model_name, filename=file, token=HUGGINGFACE_HUB_TOKEN)
                await self.upload_file(local_file, f"{model_name_s3}/{file}")

    async def load_model_from_s3(self, model_name: str):
        """Carga el modelo desde S3."""
        model_name_s3 = model_name.replace("/", "-").lower()
        files = {
            "model": f"{model_name_s3}/pytorch_model.bin",
            "tokenizer": f"{model_name_s3}/tokenizer.json",
            "config": f"{model_name_s3}/config.json",
        }

        for key, path in files.items():
            if not await self.file_exists(path):
                raise HTTPException(status_code=404, detail=f"Archivo {path} no encontrado en S3.")

        model_bytes = await self.get_file(files["model"])
        tokenizer_bytes = await self.get_file(files["tokenizer"])
        config_bytes = await self.get_file(files["config"])

        model = AutoModelForCausalLM.from_pretrained(model_bytes, config=config_bytes)
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_bytes)

        return model, tokenizer

@app.post("/generate")
async def generate(request: GenerateRequest):
    try:
        # Validaciones iniciales
        if not request.model_name or not request.input_text or not request.task_type:
            raise HTTPException(status_code=400, detail="Todos los campos son obligatorios.")

        if request.task_type not in ["text-to-text", "text-to-image", "text-to-speech", "text-to-video"]:
            raise HTTPException(status_code=400, detail="Tipo de tarea no soportado.")

        # Descarga y carga del modelo
        s3_manager = S3Manager(S3_BUCKET_NAME)
        await s3_manager.download_model_files(request.model_name)
        model, tokenizer = await s3_manager.load_model_from_s3(request.model_name)

        # Generación según el tipo de tarea
        if request.task_type == "text-to-text":
            generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
            result = generator(request.input_text, max_length=MAX_TOKENS, num_return_sequences=1)
            return {"result": result[0]["generated_text"]}

        elif request.task_type == "text-to-image":
            generator = pipeline("text-to-image", model=model, tokenizer=tokenizer)
            image = generator(request.input_text)
            return {"image": image}

        elif request.task_type == "text-to-speech":
            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer)
            audio = generator(request.input_text)
            return {"audio": audio}

        elif request.task_type == "text-to-video":
            generator = pipeline("text-to-video", model=model, tokenizer=tokenizer)
            video = generator(request.input_text)
            return {"video": video}

    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error en la generación: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)