Update app.py

app.py

@@ -1,18 +1,17 @@
 import os
+import json
 import boto3
-import torch
+import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import safetensors.torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from fastapi.responses import StreamingResponse
-import io
-import requests
-import uvicorn
+from huggingface_hub import hf_hub_download
+from io import BytesIO
+import torch
+import safetensors
 from dotenv import load_dotenv
+import tqdm
 import re
-from tqdm import tqdm
-from huggingface_hub import hf_hub_download
 
 # Cargar las variables de entorno desde el archivo .env
 load_dotenv()
@@ -54,7 +53,7 @@ class S3DirectStream:
         try:
             print(f"[INFO] Descargando archivo {key} desde S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return response['Body']  # Devolver el cuerpo directamente para el StreamingResponse
+            return BytesIO(response['Body'].read())  # Devolver el cuerpo como BytesIO
         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:
@@ -68,9 +67,11 @@ class S3DirectStream:
         except self.s3_client.exceptions.ClientError:
             return False
 
-    def load_model_from_stream(self, model_name):
+    def load_model_from_s3(self, model_name):
         try:
             print(f"[INFO] Cargando el modelo {model_name} desde S3...")
+
+            # Verificar si el modelo existe en S3
             model_prefix = model_name.lower()
             model_files = self.get_model_file_parts(model_prefix)
 
@@ -82,16 +83,16 @@ class S3DirectStream:
             if not model_files:
                 raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
-            # Leer y cargar todos los archivos del modelo
+            # Cargar todos los archivos del modelo desde S3
             model_streams = []
-            for model_file in tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
+            for model_file in tqdm.tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
                 model_streams.append(self.stream_from_s3(model_file))
 
             config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-            config_data = config_stream.read().decode("utf-8")
+            config_data = json.loads(config_stream.read().decode("utf-8"))
 
             # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
-            if model_files[0].endswith("model.safetensors"):
+            if any(file.endswith("model.safetensors") for file in model_files):
                 print("[INFO] Cargando el modelo como safetensor...")
                 model = AutoModelForCausalLM.from_config(config_data)
                 model.load_state_dict(safetensors.torch.load_stream(model_streams[0]))
@@ -107,7 +108,7 @@ class S3DirectStream:
             print(f"[ERROR] Error al cargar el modelo desde S3: {e}")
             raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.")
 
-    def load_tokenizer_from_stream(self, model_name):
+    def load_tokenizer_from_s3(self, model_name):
         try:
             print(f"[INFO] Cargando el tokenizer {model_name} desde S3...")
             tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
@@ -121,9 +122,10 @@ class S3DirectStream:
         print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_name}...")
         files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
         model_files = []
-        for obj in tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
+        for obj in tqdm.tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
             key = obj['Key']
-            if re.match(rf"{model_name}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"):
+            # Verificar si es un archivo relevante del modelo
+            if re.match(rf"{model_name}/.*", key):
                 model_files.append(key)
         if not model_files:
             print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.")
@@ -132,18 +134,16 @@ class S3DirectStream:
     def download_and_upload_from_huggingface(self, model_name):
         try:
             print(f"[INFO] Descargando {model_name} desde Hugging Face...")
-            # Descargar los archivos necesarios de Hugging Face usando huggingface_hub
-            files_to_download = [
-                "pytorch_model.bin",
-                "config.json",
-                "tokenizer.json",
-            ]
+            # Descargar todos los archivos del modelo
+            files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN)
 
+            # Subir a S3 los archivos descargados
             for file in files_to_download:
-                # Descargar cada archivo desde Hugging Face y subirlo a S3
-                file_path = hf_hub_download(repo_id=model_name, filename=file, use_auth_token=HUGGINGFACE_TOKEN)
-                s3_key = f"{model_name}/{file}"
-                self.upload_file_to_s3(file_path, s3_key)
+                file_name = os.path.basename(file)
+                s3_key = f"{model_name}/{file_name}"
+                if not self.file_exists_in_s3(s3_key):
+                    self.upload_file_to_s3(file, s3_key)
+
         except Exception as e:
             print(f"[ERROR] Error al descargar y subir modelo desde Hugging Face: {e}")
             raise HTTPException(status_code=500, detail="Error al descargar y subir modelo desde Hugging Face.")
@@ -153,44 +153,53 @@ class S3DirectStream:
             print(f"[INFO] Subiendo archivo {file_path} a S3 con key {s3_key}...")
             with open(file_path, 'rb') as data:
                 self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data)
+            # Eliminar archivo local después de ser subido
+            os.remove(file_path)
         except Exception as e:
             print(f"[ERROR] Error al subir archivo a S3: {e}")
             raise HTTPException(status_code=500, detail="Error al subir archivo a S3.")
 
+
 @app.post("/predict/")
 async def predict(model_request: DownloadModelRequest):
     try:
         print(f"[INFO] Recibiendo solicitud para predecir con el modelo {model_request.model_name}...")
         streamer = S3DirectStream(S3_BUCKET_NAME)
-        model = streamer.load_model_from_stream(model_request.model_name)
-        tokenizer = streamer.load_tokenizer_from_stream(model_request.model_name)
+        model = streamer.load_model_from_s3(model_request.model_name)
+        tokenizer = streamer.load_tokenizer_from_s3(model_request.model_name)
 
         task = model_request.pipeline_task
-        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-image", "text-to-audio", "text-to-video"]:
+        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", 
+                        "text-to-speech", "text-to-video", "text-to-image"]:
             raise HTTPException(status_code=400, detail="Pipeline task no soportado")
 
-        nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, max_length=2046)
-
-        input_text = model_request.input_text
-        print(f"[INFO] Ejecutando tarea {task} con el texto de entrada...")
-        outputs = nlp_pipeline(input_text)
-
-        # Eliminación de archivo local después de subir a S3
-        if task == "text-to-speech":
-            s3_key = f"{model_request.model_name}/generated_audio.wav"
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav")
-        elif task == "text-to-image":
-            s3_key = f"{model_request.model_name}/generated_image.png"
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
+        # Configurar el pipeline de transformers según la tarea
+        nlp_pipeline = None
+        if task == "text-generation":
+            nlp_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        elif task == "sentiment-analysis":
+            nlp_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
+        elif task == "translation":
+            nlp_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
+        elif task == "fill-mask":
+            nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)
+        elif task == "question-answering":
+            nlp_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
+        elif task == "text-to-speech":
+            nlp_pipeline = pipeline("text-to-speech", model=model, tokenizer=tokenizer)
         elif task == "text-to-video":
-            s3_key = f"{model_request.model_name}/generated_video.mp4"
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
+            nlp_pipeline = pipeline("text-to-video", model=model, tokenizer=tokenizer)
+        elif task == "text-to-image":
+            nlp_pipeline = pipeline("text-to-image", model=model, tokenizer=tokenizer)
+
+        result = nlp_pipeline(model_request.input_text)
+        return {"result": result}
 
-        return {"output": outputs}
     except Exception as e:
-        print(f"[ERROR] Error al procesar la solicitud: {str(e)}")
-        raise HTTPException(status_code=500, detail="Error al realizar la predicción.")
+        print(f"[ERROR] Error en el proceso de predicción: {str(e)}")
+        raise HTTPException(status_code=500, detail="Error en el proceso de predicción")
+
 
+# Ejecutar la app con Uvicorn
 if __name__ == "__main__":
-    print("Iniciando servidor de predicciones en localhost:8000")
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)