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)