app.py

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
)
from io import BytesIO
import boto3
from botocore.exceptions import NoCredentialsError
from huggingface_hub import snapshot_download
import shutil

# Configuración global
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")

# Diccionario global de tokens y configuraciones
token_dict = {}

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

# Modelo de solicitud
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] = []

# Clase para cargar y gestionar los modelos desde S3
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:
            # Verificar si el modelo ya está en S3
            try:
                self.s3_client.head_object(Bucket=self.bucket_name, Key=f'{model_name}/model')
                print(f"Modelo {model_name} ya existe en S3.")
            except self.s3_client.exceptions.ClientError:
                print(f"Modelo {model_name} no existe en S3. Descargando desde Hugging Face...")

                # Eliminar caché local de Hugging Face (si existe)
                local_cache_dir = os.path.join(os.getenv("HOME"), ".cache/huggingface/hub/models--")
                if os.path.exists(local_cache_dir):
                    shutil.rmtree(local_cache_dir)

                model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN)

                # Cargar el modelo y tokenizer
                model = AutoModelForCausalLM.from_pretrained(model_path)
                tokenizer = AutoTokenizer.from_pretrained(model_path)

                # Asignar EOS y PAD token si no están definidos
                if tokenizer.eos_token_id is None:
                    tokenizer.eos_token_id = tokenizer.pad_token_id

                # Guardar el modelo y el tokenizer en el diccionario
                token_dict[model_name] = {
                    "model": model,
                    "tokenizer": tokenizer,
                    "pad_token_id": tokenizer.pad_token_id,
                    "eos_token_id": tokenizer.eos_token_id
                }

                # Subir los archivos del modelo y tokenizer a S3
                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')

                # Eliminar los archivos locales después de haber subido a S3
                shutil.rmtree(model_path)

            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}")

# Instanciación del cargador de modelos
model_loader = S3ModelLoader(S3_BUCKET_NAME, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION)

# Función de generación de texto con streaming
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)

        if stop_sequences and any(stop in output_text for stop in stop_sequences):
            yield output_text
            return

# Endpoint para generar 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 tokenizer desde S3 si no existe
        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 generar 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)}")

# Ejecutar el servidor FastAPI con Uvicorn
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)