app.py

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

# Configuraciones 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 el archivo ya existe en S3
            if not await self.file_exists_in_s3(f"{model_name}/config.json"):
                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"):
                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

    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)