import os import json import uuid import requests import threading import logging from fastapi import FastAPI, HTTPException 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") 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) 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() class DownloadModelRequest(BaseModel): model_name: str pipeline_task: str input_text: str class GCSHandler: def __init__(self, bucket_name): self.bucket = storage_client.bucket(bucket_name) def file_exists(self, blob_name): 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.") 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) 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 = [ "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", ] 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" 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() if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)