Update app.py

app.py

@@ -1,15 +1,18 @@
 import os
 import json
+import uuid
+import requests
 import threading
 import logging
-from google.cloud import storage
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from pydantic import BaseModel
 from fastapi import FastAPI, HTTPException
-import requests
-import uvicorn
+from pydantic import BaseModel
+from google.cloud import storage
+from google.auth import exceptions
+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")
@@ -17,12 +20,19 @@ 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__)
 
-credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
-storage_client = storage.Client.from_service_account_info(credentials_info)
-bucket = storage_client.bucket(GCS_BUCKET_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)
+    bucket = storage_client.bucket(GCS_BUCKET_NAME)
+    logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
+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}")
 
 app = FastAPI()
 
@@ -36,120 +46,199 @@ class GCSHandler:
         self.bucket = storage_client.bucket(bucket_name)
 
     def file_exists(self, blob_name):
-        return self.bucket.blob(blob_name).exists()
-
-    def create_folder_if_not_exists(self, folder_name):
-        if not self.file_exists(folder_name):
-            self.bucket.blob(folder_name + "/").upload_from_string("")
+        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):
-        self.create_folder_if_not_exists(os.path.dirname(blob_name))
         blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(file_stream)
+        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.open("rb")
+        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)
+        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}"}
-    response = requests.get(url, headers=headers)
-    if response.status_code == 200:
+    
+    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 = [
+                "pytorch_model.bin",
+                "config.json",
+                "tokenizer.json",
+                "model.safetensors",
+            ]
+            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)
+                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}")
+
+@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",
         ]
-        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"models/{model_name}/{file_name}"
-            bucket.blob(blob_name).upload_from_string(file_content)
-    else:
-        raise HTTPException(status_code=404, detail="Error accessing Hugging Face model files.")
-
-def download_and_verify_model(model_name):
-    model_files = [
-        "pytorch_model.bin",
-        "config.json",
-        "tokenizer.json",
-        "model.safetensors",
-    ]
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    if not all(gcs_handler.file_exists(f"models/{model_name}/{file}") for file in model_files):
-        download_model_from_huggingface(model_name)
-
-def load_model_from_gcs(model_name):
-    model_files = [
-        "pytorch_model.bin",
-        "config.json",
-        "tokenizer.json",
-        "model.safetensors",
-    ]
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    model_files_streams = {
-        file: gcs_handler.download_file(f"models/{model_name}/{file}")
-        for file in model_files if gcs_handler.file_exists(f"models/{model_name}/{file}")
-    }
-    model_stream = model_files_streams.get("pytorch_model.bin") or model_files_streams.get("model.safetensors")
-    tokenizer_stream = model_files_streams.get("tokenizer.json")
-    config_stream = model_files_streams.get("config.json")
-    model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-    return model, tokenizer
-
-def load_model(model_name):
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    try:
-        return load_model_from_gcs(model_name)
-    except HTTPException:
-        download_and_verify_model(model_name)
-        return load_model_from_gcs(model_name)
-
-@app.on_event("startup")
-async def startup():
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    blobs = list(bucket.list_blobs(prefix="models/"))
-    model_names = set(blob.name.split("/")[1] for blob in blobs)
-    def download_model_thread(model_name):
-        try:
-            download_and_verify_model(model_name)
-        except Exception as e:
-            logger.error(f"Error downloading model '{model_name}': {e}")
-    threads = [threading.Thread(target=download_model_thread, args=(model_name,)) for model_name in model_names]
-    for thread in threads:
-        thread.start()
-    for thread in threads:
-        thread.join()
-
-@app.post("/predict/")
-async def predict(request: DownloadModelRequest):
-    model_name = request.model_name
-    pipeline_task = request.pipeline_task
-    input_text = request.input_text
-    model, tokenizer = load_model(model_name)
-    pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
-    result = pipe(input_text)
-    return {"result": result}
-
+        
+        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.")
+        
+        # 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)
+                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)
+                
+                # 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}")
+                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":
+            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":
+            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.")
+
+        # Tarea de generación de modelo 3D (simulada)
+        elif request.pipeline_task == "text-to-3d":
+            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}")
+
+# Función para ejecutar en segundo plano la descarga de modelos
 def download_all_models_in_background():
     models_url = "https://huggingface.co/api/models"
-    response = requests.get(models_url)
-    if response.status_code == 200:
+    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:
-            download_model_from_huggingface(model["id"])
-
+            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()
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)