import os import torch from fastapi import FastAPI, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel from transformers import ( AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline ) import asyncio from io import BytesIO from botocore.exceptions import NoCredentialsError import boto3 from huggingface_hub import snapshot_download # Diccionario global para almacenar los tokens y configuraciones de los modelos token_dict = {} # Configuración para acceso a modelos en Hugging Face o S3 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") # Inicialización de la aplicación FastAPI app = FastAPI() # Modelo de la solicitud para la API class GenerateRequest(BaseModel): model_name: str input_text: str task_type: str temperature: float = 1.0 max_new_tokens: int = 200 stream: bool = True top_p: float = 1.0 top_k: int = 50 repetition_penalty: float = 1.0 num_return_sequences: int = 1 do_sample: bool = True chunk_delay: float = 0.0 stop_sequences: list[str] = [] class S3ModelLoader: def __init__(self, bucket_name, aws_access_key_id=None, aws_secret_access_key=None, aws_region=None): self.bucket_name = 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 ) def _get_s3_uri(self, model_name): return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}" def load_model_and_tokenizer(self, model_name): if model_name in token_dict: return token_dict[model_name] s3_uri = self._get_s3_uri(model_name) try: # Descargamos el modelo y el tokenizer desde Hugging Face directamente a S3 model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN) model = AutoModelForCausalLM.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) if tokenizer.eos_token_id is None: tokenizer.eos_token_id = tokenizer.pad_token_id # Guardamos en el diccionario global token_dict[model_name] = { "model": model, "tokenizer": tokenizer, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id } # Subimos los modelos al S3 si es necesario self.s3_client.upload_file(model_path, self.bucket_name, f'{model_name}/model') self.s3_client.upload_file(f'{model_path}/tokenizer', self.bucket_name, f'{model_name}/tokenizer') return token_dict[model_name] except NoCredentialsError: raise HTTPException(status_code=500, detail="AWS credentials not found.") except Exception as e: raise HTTPException(status_code=500, detail=f"Error loading model: {e}") model_loader = S3ModelLoader(S3_BUCKET_NAME, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION) # Función para hacer streaming de texto, generando un token a la vez async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay, max_length=2048): encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device) input_length = encoded_input["input_ids"].shape[1] remaining_tokens = max_length - input_length if remaining_tokens <= 0: yield "" generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens) def stop_criteria(input_ids, scores): decoded_output = tokenizer.decode(int(input_ids[0][-1]), skip_special_tokens=True) return decoded_output in stop_sequences stopping_criteria = StoppingCriteriaList([stop_criteria]) output_text = "" outputs = model.generate( **encoded_input, do_sample=generation_config.do_sample, max_new_tokens=generation_config.max_new_tokens, temperature=generation_config.temperature, top_p=generation_config.top_p, top_k=generation_config.top_k, repetition_penalty=generation_config.repetition_penalty, num_return_sequences=generation_config.num_return_sequences, stopping_criteria=stopping_criteria, output_scores=True, return_dict_in_generate=True ) for output in outputs.sequences: for token_id in output: token = tokenizer.decode(token_id, skip_special_tokens=True) yield token await asyncio.sleep(chunk_delay) # Simula el delay entre tokens if stop_sequences and any(stop in output_text for stop in stop_sequences): yield output_text return # Endpoint para la generación de texto @app.post("/generate") async def generate(request: GenerateRequest): try: model_name = request.model_name input_text = request.input_text temperature = request.temperature max_new_tokens = request.max_new_tokens stream = request.stream top_p = request.top_p top_k = request.top_k repetition_penalty = request.repetition_penalty num_return_sequences = request.num_return_sequences do_sample = request.do_sample chunk_delay = request.chunk_delay stop_sequences = request.stop_sequences # Cargar el modelo y el tokenizer desde el S3 model_data = model_loader.load_model_and_tokenizer(model_name) model = model_data["model"] tokenizer = model_data["tokenizer"] pad_token_id = model_data["pad_token_id"] eos_token_id = model_data["eos_token_id"] device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) generation_config = GenerationConfig( temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, do_sample=do_sample, num_return_sequences=num_return_sequences, ) return StreamingResponse( stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay), media_type="text/plain" ) except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # Endpoint para la generación de imágenes @app.post("/generate-image") async def generate_image(request: GenerateRequest): try: validated_body = request device = "cuda" if torch.cuda.is_available() else "cpu" image_generator = pipeline("text-to-image", model=validated_body.model_name, device=device) image = image_generator(validated_body.input_text)[0] img_byte_arr = BytesIO() image.save(img_byte_arr, format="PNG") img_byte_arr.seek(0) return StreamingResponse(img_byte_arr, media_type="image/png") except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # Endpoint para la generación de texto a voz @app.post("/generate-text-to-speech") async def generate_text_to_speech(request: GenerateRequest): try: validated_body = request device = "cuda" if torch.cuda.is_available() else "cpu" audio_generator = pipeline("text-to-speech", model=validated_body.model_name, device=device) audio = audio_generator(validated_body.input_text)[0] audio_byte_arr = BytesIO() audio.save(audio_byte_arr) audio_byte_arr.seek(0) return StreamingResponse(audio_byte_arr, media_type="audio/wav") except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # Endpoint para la generación de video @app.post("/generate-video") async def generate_video(request: GenerateRequest): try: validated_body = request device = "cuda" if torch.cuda.is_available() else "cpu" video_generator = pipeline("text-to-video", model=validated_body.model_name, device=device) video = video_generator(validated_body.input_text)[0] video_byte_arr = BytesIO() video.save(video_byte_arr) video_byte_arr.seek(0) return StreamingResponse(video_byte_arr, media_type="video/mp4") except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # Configuración para ejecutar el servidor if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)