Update app.py

app.py

@@ -12,6 +12,7 @@ import uvicorn
 from dotenv import load_dotenv
 import re
 from tqdm import tqdm
+from huggingface_hub import hf_hub_download
 
 # Cargar las variables de entorno desde el archivo .env
 load_dotenv()
@@ -72,8 +73,12 @@ class S3DirectStream:
             print(f"[INFO] Cargando el modelo {model_name} desde S3...")
             model_prefix = model_name.lower()
             model_files = self.get_model_file_parts(model_prefix)
-            if not model_files:
-                print(f"[INFO] Modelo no encontrado en S3, descargando desde Hugging Face a S3...")
+            
+            # Verificar si el modelo está en S3
+            if model_files:
+                print(f"[INFO] Modelo encontrado en S3, cargando directamente...")
+            else:
+                print(f"[INFO] Modelo no encontrado en S3, descargando desde Hugging Face...")
                 self.download_and_upload_from_huggingface(model_name)
                 model_files = self.get_model_file_parts(model_prefix)
             
@@ -85,6 +90,7 @@ class S3DirectStream:
             config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read().decode("utf-8")
             
+            # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
             if model_files[0].endswith("model.safetensors"):
                 print("[INFO] Cargando el modelo como safetensor...")
                 model = AutoModelForCausalLM.from_config(config_data)
@@ -126,34 +132,30 @@ class S3DirectStream:
     def download_and_upload_from_huggingface(self, model_name):
         try:
             print(f"[INFO] Descargando {model_name} desde Hugging Face...")
-            model_url = f"https://huggingface.co/{model_name}/resolve/main/"
+            # Descargar los archivos necesarios de Hugging Face usando huggingface_hub
             files_to_download = [
-                f"{model_name}/pytorch_model.bin",
-                f"{model_name}/config.json",
-                f"{model_name}/tokenizer.json",
+                "pytorch_model.bin",
+                "config.json",
+                "tokenizer.json",
             ]
-            
+
             for file in files_to_download:
-                file_url = model_url + file
-                s3_key = file
-                self.download_and_upload_to_s3_url(file_url, s3_key)
+                # 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)
         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.")
 
-    def download_and_upload_to_s3_url(self, url: str, s3_key: str):
+    def upload_file_to_s3(self, file_path, s3_key):
         try:
-            print(f"[INFO] Descargando archivo desde {url}...")
-            response = requests.get(url)
-            if response.status_code == 200:
-                print(f"[INFO] Subiendo archivo a S3 con key {s3_key}...")
-                self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=response.content)
-            else:
-                print(f"[ERROR] Error al descargar el archivo desde {url}, código de estado {response.status_code}.")
-                raise HTTPException(status_code=500, detail=f"Error al descargar el archivo desde {url}")
+            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)
         except Exception as e:
-            print(f"[ERROR] Error al procesar la URL {url}: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Error al procesar la URL {url}")
+            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):
@@ -184,13 +186,11 @@ async def predict(model_request: DownloadModelRequest):
             s3_key = f"{model_request.model_name}/generated_video.mp4"
             return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
 
-        return {"input_text": input_text, "output": outputs}
-
+        return {"output": outputs}
     except Exception as e:
-        print(f"[ERROR] Error en la predicción: {str(e)}")
+        print(f"[ERROR] Error en el proceso de predicción: {str(e)}")
         raise HTTPException(status_code=500, detail="Error al realizar la predicción.")
 
-# Iniciar servidor de predicciones
 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)