import os import json from fastapi import FastAPI, HTTPException from pydantic import BaseModel from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline import boto3 import logging from huggingface_hub import hf_hub_download # 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") HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") # Cliente de 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 ) app = FastAPI() class GenerateRequest(BaseModel): model_name: str input_text: str task_type: 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 async def download_and_upload_to_s3(self, model_name): try: model_name = model_name.replace("/", "-").lower() # Descargar el archivo config.json desde Hugging Face config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN) tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN) # Verificar si la carpeta y los archivos ya existen en S3 if not await self.file_exists_in_s3(f"{model_name}/config.json"): logging.info(f"El archivo config.json no existe en S3. Subiendo desde Hugging Face...") self.create_folder_if_not_exists(model_name) with open(config_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file) if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"): logging.info(f"El archivo tokenizer.json no existe en S3. Subiendo desde Hugging Face...") self.create_folder_if_not_exists(model_name) with open(tokenizer_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file) except Exception as e: logging.error(f"Error al cargar el modelo desde Hugging Face a S3: {e}") raise HTTPException(status_code=500, detail=f"Error al cargar el modelo: {str(e)}") async def file_exists_in_s3(self, s3_key): try: self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key) return True except self.s3_client.exceptions.ClientError: return False def create_folder_if_not_exists(self, model_name): try: # Las carpetas no existen como tal en S3, pero se pueden crear archivos vacíos para simular carpetas # Crear un archivo vacío para simular la carpeta self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/") except Exception as e: logging.error(f"Error al crear la carpeta en S3: {e}") raise HTTPException(status_code=500, detail=f"Error al crear la carpeta en S3: {str(e)}") async def load_model_from_s3(self, model_name): try: model_name = model_name.replace("/", "-").lower() model_files = await self.get_model_file_parts(model_name) if not model_files: await self.download_and_upload_to_s3(model_name) # Cargar configuración del modelo desde S3 config_data = await self.stream_from_s3(f"{model_name}/config.json") if isinstance(config_data, bytes): config_data = config_data.decode("utf-8") config_json = json.loads(config_data) # Cargar el modelo model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json) return model except HTTPException as e: raise e except Exception as e: logging.error(f"Error al cargar el modelo desde S3: {e}") raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {str(e)}") async def load_tokenizer_from_s3(self, model_name): try: model_name = model_name.replace("/", "-").lower() tokenizer_data = await self.stream_from_s3(f"{model_name}/tokenizer.json") if isinstance(tokenizer_data, bytes): tokenizer_data = tokenizer_data.decode("utf-8") tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}") return tokenizer except Exception as e: logging.error(f"Error al cargar el tokenizer desde S3: {e}") raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {str(e)}") async def stream_from_s3(self, key): 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"El archivo {key} no existe en el bucket S3.") except Exception as e: raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}") async def get_model_file_parts(self, model_name): try: model_name = model_name.replace("/", "-").lower() files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name) model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']] return model_files except Exception as e: raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {str(e)}") @app.post("/generate") async def generate(request: GenerateRequest): try: model_name = request.model_name input_text = request.input_text task_type = request.task_type s3_direct_stream = S3DirectStream(S3_BUCKET_NAME) model = await s3_direct_stream.load_model_from_s3(model_name) tokenizer = await s3_direct_stream.load_tokenizer_from_s3(model_name) if task_type == "text-to-text": generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0) result = generator(input_text, max_length=1024, num_return_sequences=1) return {"result": result[0]["generated_text"]} elif task_type == "text-to-image": generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=0) image = generator(input_text) return {"result": image} elif task_type == "text-to-speech": generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=0) audio = generator(input_text) return {"result": audio} elif task_type == "text-to-video": generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=0) video = generator(input_text) return {"result": video} else: raise HTTPException(status_code=400, detail="Tipo de tarea no soportada") except HTTPException as e: raise e except Exception as e: raise HTTPException(status_code=500, detail=str(e)) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)