Update app.py

app.py

@@ -1,14 +1,12 @@
 import os
-import logging
-import boto3
+import json
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from huggingface_hub import hf_hub_download
-import torch
-import safetensors
+import boto3
+import logging
 import asyncio
-from tqdm import tqdm  # Importar tqdm para la barra de progreso
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -55,14 +53,11 @@ class S3DirectStream:
 
     def _stream_from_s3(self, key):
         try:
-            logger.info(f"Descargando archivo {key} desde S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            file_content = response['Body'].read()  # This returns a bytes object
-            return file_content
+            return response['Body'].read()  # This is a bytes object
         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:
-            logger.error(f"Error al descargar {key} desde S3: {str(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):
@@ -72,48 +67,71 @@ class S3DirectStream:
     def _get_model_file_parts(self, model_name):
         try:
             model_name = model_name.replace("/", "-").lower()
-            logger.info(f"Obteniendo archivos del modelo {model_name} desde S3...")
             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 obtener archivos del modelo {model_name} desde S3: {e}")
 
     async def load_model_from_s3(self, model_name):
         try:
-            logger.info(f"Cargando modelo {model_name} desde S3...")
             model_name = model_name.replace("/", "-").lower()
             model_files = await self.get_model_file_parts(model_name)
 
-            if 'pytorch_model.bin' not in model_files:
-                raise HTTPException(status_code=404, detail="Archivo 'pytorch_model.bin' no encontrado en S3")
-            if 'tokenizer.json' not in model_files:
-                raise HTTPException(status_code=404, detail="Archivo 'tokenizer.json' no encontrado en S3")
+            if not model_files:
+                await self.download_and_upload_to_s3(model_name)
 
-            model_bytes = await self.stream_from_s3(f"{model_name}/pytorch_model.bin")
-            logger.info(f"Modelo descargado correctamente. Cargando el modelo en memoria...")
-            model = AutoModelForCausalLM.from_pretrained(model_bytes, config=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:
-            logger.error(f"Error al cargar el modelo desde S3: {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:
-            logger.info(f"Cargando tokenizer del modelo {model_name} desde S3...")
             model_name = model_name.replace("/", "-").lower()
-            tokenizer_bytes = await self.stream_from_s3(f"{model_name}/tokenizer.json")
-            tokenizer = AutoTokenizer.from_pretrained(tokenizer_bytes)
+            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:
-            logger.error(f"Error al cargar el tokenizer desde S3: {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):
         try:
             if force_download:
@@ -121,84 +139,84 @@ class S3DirectStream:
             
             model_name = model_name.replace("/", "-").lower()
 
-            if not await self.file_exists_in_s3(f"{model_name}/pytorch_model.bin") or not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                logger.info(f"Descargando archivos del modelo {model_name} desde Hugging Face...")
-                model_file = hf_hub_download(repo_id=model_name, filename="pytorch_model.bin", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
+            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}/pytorch_model.bin"):
-                    with open(model_file, "rb") as file:
-                        logger.info(f"Cargando archivo {model_name}/pytorch_model.bin a S3...")
-                        self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/pytorch_model.bin", Body=file)
+                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:
-                        logger.info(f"Cargando archivo {model_name}/tokenizer.json a S3...")
                         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:
-            logger.error(f"Error al descargar o cargar archivos desde Hugging Face a S3: {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)
+
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga del modelo: {e}")
 
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
-        task_type = request.task_type
         model_name = request.model_name
         input_text = request.input_text
-
-        logger.info(f"Iniciando la generación para el modelo {model_name} con el tipo de tarea {task_type}...")
-
+        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)
-
-        logger.info(f"Modelo y tokenizer cargados correctamente. Procesando tarea {task_type}...")
-
+        
+        # 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)
-            logger.info(f"Generación completada: {result[0]['generated_text']}")
             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)
-            logger.info(f"Imagen generada.")
-            return {"image": image}
-
-        elif task_type == "text-to-video":
-            generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=0)
-            video = generator(input_text)
-            logger.info(f"Video generado.")
-            return {"video": video}
+            return {"result": image}
 
         elif task_type == "text-to-speech":
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=0)
             audio = generator(input_text)
-            logger.info(f"Audio generado.")
-            return {"audio": audio}
+            return {"result": audio}
 
-        elif task_type == "text-to-audio":
-            generator = pipeline("text-to-audio", model=model, tokenizer=tokenizer, device=0)
-            audio = generator(input_text)
-            logger.info(f"Audio generado.")
-            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 {"result": video}
 
         else:
-            raise HTTPException(status_code=400, detail="Tipo de tarea no soportado.")
+            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=f"Error en la generación: {str(e)}")
-
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn