Update app.py

app.py

@@ -12,7 +12,6 @@ from transformers import pipeline
 from dotenv import load_dotenv
 import uvicorn
 
-# Configuración de carga de variables de entorno
 load_dotenv()
 
 API_KEY = os.getenv("API_KEY")
@@ -20,11 +19,9 @@ GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
 GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
-# Configuración del logger
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Inicializar el cliente de Google Cloud Storage
 try:
     credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
     storage_client = storage.Client.from_service_account_info(credentials_info)
@@ -65,7 +62,7 @@ class GCSHandler:
             logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
             raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
         logger.debug(f"Descargando archivo '{blob_name}' de GCS.")
-        return blob.open("rb")  # Abre el archivo en modo lectura de bytes
+        return blob.open("rb")
 
     def generate_signed_url(self, blob_name, expiration=3600):
         blob = self.bucket.blob(blob_name)
@@ -76,7 +73,6 @@ class GCSHandler:
 def download_model_from_huggingface(model_name):
     url = f"https://huggingface.co/{model_name}/tree/main"
     headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    
     try:
         logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
         response = requests.get(url, headers=headers)
@@ -90,7 +86,6 @@ def download_model_from_huggingface(model_name):
             for file_name in model_files:
                 file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
                 file_content = requests.get(file_url).content
-                # Subir el archivo directamente desde el contenido
                 blob_name = f"{model_name}/{file_name}"
                 blob = bucket.blob(blob_name)
                 blob.upload_from_string(file_content)
@@ -131,14 +126,12 @@ async def predict(request: DownloadModelRequest):
             logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
             raise HTTPException(status_code=500, detail="Required model files missing.")
         
-        # Tareas basadas en texto
         if request.pipeline_task in ["text-generation", "translation", "summarization"]:
             pipe = pipeline(request.pipeline_task, model=model_stream, tokenizer=tokenizer_stream)
             result = pipe(request.input_text)
             logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
             return {"response": result[0]}
 
-        # Tareas de imagen
         elif request.pipeline_task == "image-generation":
             try:
                 pipe = pipeline("image-generation", model=model_stream)
@@ -148,7 +141,6 @@ async def predict(request: DownloadModelRequest):
                 image_path = f"images/{image_filename}"
                 image.save(image_path)
                 
-                # Subir la imagen generada a GCS
                 gcs_handler.upload_file(image_path, open(image_path, "rb"))
                 image_url = gcs_handler.generate_signed_url(image_path)
                 logger.info(f"Imagen generada y subida correctamente con URL: {image_url}")
@@ -189,7 +181,6 @@ async def predict(request: DownloadModelRequest):
                 logger.error(f"Error transformando la imagen: {e}")
                 raise HTTPException(status_code=400, detail="Error transformando la imagen.")
 
-        # Tarea de generación de modelo 3D (simulada)
         elif request.pipeline_task == "text-to-3d":
             try:
                 model_3d_filename = f"{uuid.uuid4().hex}.obj"
@@ -199,46 +190,16 @@ async def predict(request: DownloadModelRequest):
 
                 gcs_handler.upload_file(f"3d-models/{model_3d_filename}", open(model_3d_path, "rb"))
                 model_3d_url = gcs_handler.generate_signed_url(f"3d-models/{model_3d_filename}")
-                logger.info(f"Modelo 3D generado y subido con URL: {model_3d_url}")
+                logger.info(f"Modelo 3D generado y subido correctamente con URL: {model_3d_url}")
                 return {"response": {"model_3d_url": model_3d_url}}
+
             except Exception as e:
                 logger.error(f"Error generando el modelo 3D: {e}")
                 raise HTTPException(status_code=400, detail="Error generando el modelo 3D.")
-            
-    except HTTPException as e:
-        logger.error(f"HTTPException: {e.detail}")
-        raise e
-    except Exception as e:
-        logger.error(f"Error inesperado: {e}")
-        raise HTTPException(status_code=500, detail=f"Error: {e}")
 
-# Función para ejecutar en segundo plano la descarga de modelos
-def download_all_models_in_background():
-    models_url = "https://huggingface.co/api/models"
-    try:
-        logger.info("Obteniendo lista de modelos desde Hugging Face...")
-        response = requests.get(models_url)
-        if response.status_code != 200:
-            logger.error("Error al obtener la lista de modelos de Hugging Face.")
-            raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
-        
-        models = response.json()
-        for model in models:
-            model_name = model["id"]
-            logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
-            download_model_from_huggingface(model_name)
     except Exception as e:
-        logger.error(f"Error al descargar modelos en segundo plano: {e}")
-        raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
-
-# Iniciar la descarga de modelos en segundo plano
-def run_in_background():
-    logger.info("Iniciando la descarga de modelos en segundo plano...")
-    threading.Thread(target=download_all_models_in_background, daemon=True).start()
-
-@app.on_event("startup")
-async def startup_event():
-    run_in_background()
+        logger.error(f"Error en la predicción: {e}")
+        raise HTTPException(status_code=500, detail=f"Error en la predicción: {e}")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)