Update app.py

app.py

@@ -1,18 +1,17 @@
 import os
 import json
-import logging
 import uuid
+import requests
 import threading
-import io
+import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import uvicorn
-import torch
-import requests
-from safetensors import safe_open
+from google.auth import exceptions
+from transformers import pipeline
 from dotenv import load_dotenv
+import uvicorn
+import io
 
 load_dotenv()
 
@@ -29,7 +28,7 @@ try:
     storage_client = storage.Client.from_service_account_info(credentials_info)
     bucket = storage_client.bucket(GCS_BUCKET_NAME)
     logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
-except (json.JSONDecodeError, KeyError, ValueError) as e:
+except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
     logger.error(f"Error al cargar las credenciales o bucket: {e}")
     raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
 
@@ -45,60 +44,39 @@ class GCSHandler:
         self.bucket = storage_client.bucket(bucket_name)
 
     def file_exists(self, blob_name):
-        return self.bucket.blob(blob_name).exists()
+        exists = self.bucket.blob(blob_name).exists()
+        logger.debug(f"Comprobando existencia de archivo '{blob_name}': {exists}")
+        return exists
+
+    def upload_file(self, blob_name, file_stream):
+        blob = self.bucket.blob(blob_name)
+        try:
+            blob.upload_from_file(file_stream)
+            logger.info(f"Archivo '{blob_name}' subido exitosamente a GCS.")
+        except Exception as e:
+            logger.error(f"Error subiendo el archivo '{blob_name}' a GCS: {e}")
+            raise HTTPException(status_code=500, detail=f"Error subiendo archivo '{blob_name}' a GCS")
 
     def download_file(self, blob_name):
         blob = self.bucket.blob(blob_name)
         if not blob.exists():
+            logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
             raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
-        return blob.download_as_bytes()
-
-    def upload_file(self, blob_name, file_data):
-        blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(file_data)
+        logger.debug(f"Descargando archivo '{blob_name}' de GCS.")
+        return blob.open("rb")  # Abre el archivo en modo lectura de bytes
 
     def generate_signed_url(self, blob_name, expiration=3600):
         blob = self.bucket.blob(blob_name)
-        return blob.generate_signed_url(expiration=expiration)
-
-    def create_folder(self, folder_name):
-        blob = self.bucket.blob(folder_name + "/")
-        blob.upload_from_string("")  # Create an empty "folder"
-
-def load_model_from_gcs(model_name: str, model_files: list):
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    model_blobs = {file: gcs_handler.download_file(f"{model_name}/{file}") for file in model_files}
-    
-    model_stream = model_blobs.get("pytorch_model.bin") or model_blobs.get("model.safetensors")
-    config_stream = model_blobs.get("config.json")
-    tokenizer_stream = model_blobs.get("tokenizer.json")
-    
-    if model_stream and model_stream.endswith(".safetensors"):
-        model = load_safetensors_model(model_stream)
-    else:
-        model = AutoModelForCausalLM.from_pretrained(io.BytesIO(model_stream), config=config_stream)
-    
-    tokenizer = AutoTokenizer.from_pretrained(io.BytesIO(tokenizer_stream))
-    
-    return model, tokenizer
-
-def load_safetensors_model(model_stream):
-    with safe_open(io.BytesIO(model_stream), framework="pt") as model_data:
-        model = torch.load(model_data)
-    return model
-
-def get_model_files_from_gcs(model_name: str):
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    blob_list = list(gcs_handler.bucket.list_blobs(prefix=f"{model_name}/"))
-    model_files = [blob.name for blob in blob_list if any(part in blob.name for part in ["pytorch_model", "model"]) and "index" not in blob.name]
-    model_files = sorted(model_files)
-    return model_files
+        url = blob.generate_signed_url(expiration=expiration)
+        logger.debug(f"Generada URL firmada para '{blob_name}': {url}")
+        return url
 
 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)
         if response.status_code == 200:
             model_files = [
@@ -107,90 +85,151 @@ def download_model_from_huggingface(model_name):
                 "tokenizer.json",
                 "model.safetensors",
             ]
-            def download_file(file_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
                 blob_name = f"{model_name}/{file_name}"
                 blob = bucket.blob(blob_name)
                 blob.upload_from_string(file_content)
-            
-            threads = [threading.Thread(target=download_file, args=(file_name,)) for file_name in model_files]
-            for thread in threads:
-                thread.start()
-            for thread in threads:
-                thread.join()
+                logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
         else:
+            logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
             raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
     except Exception as e:
+        logger.error(f"Error descargando archivos de Hugging Face: {e}")
         raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
 
-def download_model_files(model_name: str):
-    model_files = get_model_files_from_gcs(model_name)
-    if not model_files:
-        download_model_from_huggingface(model_name)
-        model_files = get_model_files_from_gcs(model_name)
-    return model_files
-
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
+    logger.info(f"Iniciando predicción para el modelo '{request.model_name}' con tarea '{request.pipeline_task}'...")
     try:
         gcs_handler = GCSHandler(GCS_BUCKET_NAME)
         model_prefix = request.model_name
+        model_files = [
+            "pytorch_model.bin",
+            "config.json",
+            "tokenizer.json",
+            "model.safetensors",
+        ]
         
-        model_files = download_model_files(model_prefix)
-        model, tokenizer = load_model_from_gcs(model_prefix, model_files)
+        model_files_exist = all(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files)
+        
+        if not model_files_exist:
+            logger.info(f"Modelos no encontrados en GCS, descargando '{model_prefix}' desde Hugging Face...")
+            download_model_from_huggingface(model_prefix)
+        
+        model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
+        
+        config_stream = model_files_streams.get("config.json")
+        tokenizer_stream = model_files_streams.get("tokenizer.json")
+        model_stream = model_files_streams.get("pytorch_model.bin")
+        
+        if not config_stream or not tokenizer_stream or not model_stream:
+            logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
+            raise HTTPException(status_code=500, detail="Required model files missing.")
         
-        pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
-
         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]}
 
         elif request.pipeline_task == "image-generation":
-            images = pipe(request.input_text)
-            image = images[0]
-            image_filename = f"{uuid.uuid4().hex}.png"
-            image_path = f"images/{image_filename}"
-            image.save(image_path)
-            gcs_handler.upload_file(image_path, open(image_path, "rb"))
-            image_url = gcs_handler.generate_signed_url(image_path)
-            return {"response": {"image_url": image_url}}
+            try:
+                pipe = pipeline("image-generation", model=model_stream)
+                images = pipe(request.input_text)
+                image = images[0]
+                image_filename = f"{uuid.uuid4().hex}.png"
+                image_path = f"images/{image_filename}"
+                image.save(image_path)
+                
+                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}")
+                return {"response": {"image_url": image_url}}
+            except Exception as e:
+                logger.error(f"Error generando la imagen: {e}")
+                raise HTTPException(status_code=400, detail="Error generando la imagen.")
 
         elif request.pipeline_task == "image-editing":
-            edited_images = pipe(request.input_text)
-            edited_image = edited_images[0]
-            edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
-            edited_image.save(edited_image_filename)
-            gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_filename, "rb"))
-            edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
-            return {"response": {"edited_image_url": edited_image_url}}
+            try:
+                pipe = pipeline("image-editing", model=model_stream)
+                edited_images = pipe(request.input_text)
+                edited_image = edited_images[0]
+                edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
+                edited_image.save(edited_image_filename)
+
+                gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_filename, "rb"))
+                edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
+                logger.info(f"Imagen editada y subida correctamente con URL: {edited_image_url}")
+                return {"response": {"edited_image_url": edited_image_url}}
+            except Exception as e:
+                logger.error(f"Error editando la imagen: {e}")
+                raise HTTPException(status_code=400, detail="Error editando la imagen.")
 
         elif request.pipeline_task == "image-to-image":
-            transformed_images = pipe(request.input_text)
-            transformed_image = transformed_images[0]
-            transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
-            transformed_image.save(transformed_image_filename)
-            gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_filename, "rb"))
-            transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
-            return {"response": {"transformed_image_url": transformed_image_url}}
+            try:
+                pipe = pipeline("image-to-image", model=model_stream)
+                transformed_images = pipe(request.input_text)
+                transformed_image = transformed_images[0]
+                transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
+                transformed_image.save(transformed_image_filename)
+
+                gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_filename, "rb"))
+                transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
+                logger.info(f"Imagen transformada y subida correctamente con URL: {transformed_image_url}")
+                return {"response": {"transformed_image_url": transformed_image_url}}
+            except Exception as e:
+                logger.error(f"Error transformando la imagen: {e}")
+                raise HTTPException(status_code=400, detail="Error transformando la imagen.")
 
         elif request.pipeline_task == "text-to-3d":
-            model_3d_filename = f"{uuid.uuid4().hex}.obj"
-            model_3d_path = f"3d-models/{model_3d_filename}"
-            with open(model_3d_path, "w") as f:
-                f.write("Simulated 3D model data")
-            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}")
-            return {"response": {"model_3d_url": model_3d_url}}
+            try:
+                model_3d_filename = f"{uuid.uuid4().hex}.obj"
+                model_3d_path = f"3d-models/{model_3d_filename}"
+                with open(model_3d_path, "w") as f:
+                    f.write("Simulated 3D model data")
+
+                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}")
+                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}")
 
+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.")
+
+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():
-    logger.info("Iniciando la API...")
+    run_in_background()
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)