Update app.py

app.py

@@ -1,20 +1,13 @@
 import os
-import json
+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 boto3
-import logging
 import asyncio
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-console_handler = logging.StreamHandler()
-formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-console_handler.setFormatter(formatter)
-logger.addHandler(console_handler)
-
+# 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")
@@ -23,6 +16,7 @@ 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,
@@ -30,194 +24,135 @@ s3_client = boto3.client(
     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  # Added task type to handle different tasks (e.g., text-to-image, text-to-speech)
+    task_type: str
 
-class S3DirectStream:
+# Clase para manejo de S3
+class S3Manager:
     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
+        self.s3_client = s3_client
 
-    async def stream_from_s3(self, key):
+    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._stream_from_s3, key)
+        return await loop.run_in_executor(None, self._get_file_sync, key)
 
-    def _stream_from_s3(self, 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()  # This is a bytes object
+            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.")
+            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 descargar {key} desde S3: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Error al obtener el archivo {key} de S3: {str(e)}")
 
-    async def get_model_file_parts(self, model_name):
+    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._get_model_file_parts, model_name)
+        return await loop.run_in_executor(None, self._upload_file_sync, file_path, key)
 
-    def _get_model_file_parts(self, model_name):
+    def _upload_file_sync(self, file_path: str, key: str):
         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
+            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 obtener archivos del modelo {model_name} desde S3: {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)
-
-            config_data = await self.stream_from_s3(f"{model_name}/config.json")
-            if not config_data:
-                raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_name}/config.json está vacío o no se pudo leer.")
-            
-            # Ensure config_data is a string or bytes-like object
-            if isinstance(config_data, bytes):
-                config_data = config_data.decode("utf-8")  # Decodificar los bytes a string si es necesario
-            
-            config_json = json.loads(config_data)  # Ahora podemos usar json.loads sin problema
-
-            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False)
-            return model
-
-        except HTTPException as e:
-            raise e
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {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")
-            
-            # Ensure tokenizer_data is a string or bytes-like object
-            if isinstance(tokenizer_data, bytes):
-                tokenizer_data = tokenizer_data.decode("utf-8")  # Decodificar los bytes a string si es necesario
-            
-            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
-            return tokenizer
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}")
-
-    async def create_s3_folders(self, s3_key):
-        try:
-            folder_keys = s3_key.split('-')
-            for i in range(1, len(folder_keys)):
-                folder_key = '-'.join(folder_keys[:i]) + '/'
-                if not await self.file_exists_in_s3(folder_key):
-                    logger.info(f"Creando carpeta en S3: {folder_key}")
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=folder_key, Body='')
+            raise HTTPException(status_code=500, detail=f"Error al subir {key} a S3: {str(e)}")
 
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {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)
 
-    async def file_exists_in_s3(self, s3_key):
+    def _file_exists_sync(self, key: str):
         try:
-            self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key)
+            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_and_upload_to_s3(self, model_name, force_download=False):
-        try:
-            if force_download:
-                logger.info(f"Forzando la descarga del modelo {model_name} y la carga a S3.")
-            
-            model_name = model_name.replace("/", "-").lower()
-
-            if not await self.file_exists_in_s3(f"{model_name}/config.json") or not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
-                tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
-
-                await self.create_s3_folders(f"{model_name}/")
-
-                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)
+    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"]
 
-                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)
-            else:
-                logger.info(f"Los archivos del modelo {model_name} ya existen en S3. No es necesario descargarlos de nuevo.")
+        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}")
 
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar o cargar archivos desde Hugging Face a S3: {e}")
+    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",
+        }
 
-    async def resume_download(self, model_name):
-        try:
-            logger.info(f"Reanudando la descarga del modelo {model_name} desde Hugging Face.")
-            config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
-            tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
+        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.")
 
-            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)
+        model_bytes = await self.get_file(files["model"])
+        tokenizer_bytes = await self.get_file(files["tokenizer"])
+        config_bytes = await self.get_file(files["config"])
 
-            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)
+        model = AutoModelForCausalLM.from_pretrained(model_bytes, config=config_bytes)
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_bytes)
 
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga del modelo: {e}")
+        return model, tokenizer
 
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
-        model_name = request.model_name
-        input_text = request.input_text
-        task_type = request.task_type
-        
-        # Create an instance of S3DirectStream
-        s3_direct_stream = S3DirectStream(S3_BUCKET_NAME)
-        
-        # Load model and tokenizer
-        model = await s3_direct_stream.load_model_from_s3(model_name)
-        tokenizer = await s3_direct_stream.load_tokenizer_from_s3(model_name)
-        
-        # Generate based on task type
-        if task_type == "text-to-text":
-            generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
-            result = generator(input_text, max_length=MAX_TOKENS, num_return_sequences=1)
+        # 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 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 request.task_type == "text-to-image":
+            generator = pipeline("text-to-image", model=model, tokenizer=tokenizer)
+            image = generator(request.input_text)
+            return {"image": image}
 
-        elif task_type == "text-to-video":
-            generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=0)
-            video = generator(input_text)
-            return {"result": video}
+        elif request.task_type == "text-to-speech":
+            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer)
+            audio = generator(request.input_text)
+            return {"audio": audio}
 
-        else:
-            raise HTTPException(status_code=400, detail="Tipo de tarea no soportada")
+        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=str(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)
+