Update app.py

app.py

@@ -1,13 +1,14 @@
 import os
-import json
 import logging
 import boto3
-from fastapi import FastAPI, HTTPException, Query
-from fastapi.responses import JSONResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import hf_hub_download
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from safetensors.torch import load_file
+import torch
 import asyncio
 
+# Configuración de logs
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 console_handler = logging.StreamHandler()
@@ -15,6 +16,7 @@ formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)
 
+# Configuración de AWS y 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")
@@ -30,8 +32,16 @@ s3_client = boto3.client(
     region_name=AWS_REGION
 )
 
+# Crear la aplicación FastAPI
 app = FastAPI()
 
+# Modelo de datos para la solicitud
+class GenerateRequest(BaseModel):
+    model_name: str
+    input_text: str
+    task_type: str
+
+# Clase para gestionar el acceso a S3
 class S3DirectStream:
     def __init__(self, bucket_name):
         self.s3_client = boto3.client(
@@ -42,6 +52,7 @@ class S3DirectStream:
         )
         self.bucket_name = bucket_name
 
+    # Función para obtener el archivo desde S3
     async def stream_from_s3(self, key):
         loop = asyncio.get_event_loop()
         return await loop.run_in_executor(None, self._stream_from_s3, key)
@@ -49,155 +60,97 @@ class S3DirectStream:
     def _stream_from_s3(self, key):
         try:
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return response['Body']
+            file_content = response['Body'].read()
+            if not file_content:
+                raise HTTPException(status_code=404, detail=f"El archivo {key} está vacío.")
+            return file_content
         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):
-        loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(None, self._get_model_file_parts, model_name)
-
-    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: {e}")
-
+    # Cargar el modelo directamente desde S3
     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_stream = await self.stream_from_s3(f"{model_name}/config.json")
-            config_data = config_stream.read()
-
-            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.")
-            
-            config_text = config_data.decode("utf-8")
-            config_json = json.loads(config_text)
-
-            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False)
+            model_bytes = await self.stream_from_s3(f"{model_name}/pytorch_model.bin")
+            if model_bytes:
+                model = load_file(model_bytes)
+                return model
+            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}/pytorch_model.bin")
             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}")
 
+    # Cargar el tokenizer desde S3
     async def load_tokenizer_from_s3(self, model_name):
         try:
             model_name = model_name.replace("/", "-").lower()
-            tokenizer_stream = await self.stream_from_s3(f"{model_name}/tokenizer.json")
-            tokenizer_data = tokenizer_stream.read().decode("utf-8")
-
-            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
+            tokenizer_bytes = await self.stream_from_s3(f"{model_name}/tokenizer.json")
+            if not tokenizer_bytes:
+                raise HTTPException(status_code=404, detail="El archivo tokenizer.json está vacío o no existe.")
+            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}/tokenizer.json")
             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='')
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {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 download_and_upload_to_s3(self, model_name, force_download=False):
+    # Obtener los archivos del modelo desde S3
+    async def get_model_file_parts(self, model_name):
         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)
-
-                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.")
-
-        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 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)
-
-            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)
-
+            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']]
+            if not model_files:
+                raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados.")
+            return model_files
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga o cargar archivos desde Hugging Face a S3: {e}")
-
-def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS):
-    tokens = tokenizer.encode(text)
-    chunks = []
-    for i in range(0, len(tokens), max_tokens):
-        chunk = tokens[i:i+max_tokens]
-        chunks.append(tokenizer.decode(chunk))
-    return chunks
-
-def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
-    generated_text = ""
-    while len(input_text) > 0:
-        chunks = split_text_by_tokens(input_text, tokenizer, max_tokens)
-        for chunk in chunks:
-            generated_text += model.generate(chunk)
-    return generated_text
+            raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
 
+# Endpoint para la generación
 @app.post("/generate")
-async def generate_text(model_name: str = Query(...), input_text: str = Query(...)):
+async def generate(request: GenerateRequest):
     try:
-        model_loader = S3DirectStream(S3_BUCKET_NAME)
-        model = await model_loader.load_model_from_s3(model_name)
-        tokenizer = await model_loader.load_tokenizer_from_s3(model_name)
-        
-        chunks = split_text_by_tokens(input_text, tokenizer, max_tokens=MAX_TOKENS)
-        
-        generated_text = continue_generation(input_text, model, tokenizer)
-        
-        return {"generated_text": generated_text}
-
+        task_type = request.task_type
+        model_name = request.model_name
+        input_text = request.input_text
+
+        s3_direct_stream = S3DirectStream(S3_BUCKET_NAME)
+
+        # Cargar el modelo y tokenizer desde S3
+        model = await s3_direct_stream.load_model_from_s3(model_name)
+        tokenizer = await s3_direct_stream.load_tokenizer_from_s3(model_name)
+
+        # Generar dependiendo del tipo de tarea
+        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)
+            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 {"image": image}
+
+        elif task_type == "text-to-audio" or task_type == "text-to-speech":
+            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=0)
+            audio = generator(input_text)
+            return {"audio": audio}
+
+        elif task_type == "text-to-video":
+            generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=0)
+            video = generator(input_text)
+            return {"video": video}
+
+        else:
+            raise HTTPException(status_code=400, detail="Tipo de tarea no soportado.")
+    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)}")
 
+# Ejecutar la aplicación
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)