Update app.py

app.py

@@ -1,18 +1,24 @@
 import os
 import json
-import uuid
-import requests
-import threading
 import logging
+import io
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
 from google.auth import exceptions
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from dotenv import load_dotenv
+import torch
+import safetensors.torch
+import requests
+from diffusers import StableDiffusionPipeline
+from audiocraft.models import AudioLM
+import asyncio
+import threading
 import uvicorn
+from transformers import pipeline as tts_pipeline
+import soundfile as sf  # Para manejar el audio de salida
 
-# Configuración de carga de variables de entorno
 load_dotenv()
 
 API_KEY = os.getenv("API_KEY")
@@ -20,11 +26,10 @@ 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
+# Configuración de GCS
 try:
     credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
     storage_client = storage.Client.from_service_account_info(credentials_info)
@@ -50,6 +55,14 @@ class GCSHandler:
         logger.debug(f"Comprobando existencia de archivo '{blob_name}': {exists}")
         return exists
 
+    def download_file_as_stream(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")  # Devuelve un stream (modo lectura binaria)
+
     def upload_file(self, blob_name, file_stream):
         blob = self.bucket.blob(blob_name)
         try:
@@ -59,14 +72,6 @@ class GCSHandler:
             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)
@@ -86,6 +91,7 @@ def download_model_from_huggingface(model_name):
                 "config.json",
                 "tokenizer.json",
                 "model.safetensors",
+                "pytorch_model.bin"
             ]
             for file_name in model_files:
                 file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
@@ -102,109 +108,98 @@ def download_model_from_huggingface(model_name):
         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 load_model_from_gcs(model_name, gcs_handler):
+    model_files = {
+        "config": f"{model_name}/config.json",
+        "tokenizer": f"{model_name}/tokenizer.json",
+        "model_bin": f"{model_name}/pytorch_model.bin",
+        "model_safetensors": f"{model_name}/model.safetensors"
+    }
+
+    model_data = {}
+    for key, blob_name in model_files.items():
+        if not gcs_handler.file_exists(blob_name):
+            logger.info(f"{key.capitalize()} no encontrado en GCS, descargando desde Hugging Face...")
+            download_model_from_huggingface(model_name)
+        model_data[key] = gcs_handler.download_file_as_stream(blob_name)
+
+    return model_data
+
+def load_diffuser_model_from_streams(model_data, model_name):
+    model_bin_stream = model_data.get("model_bin")
+    model_safetensors_stream = model_data.get("model_safetensors")
+
+    if model_bin_stream or model_safetensors_stream:
+        # Cargar el modelo de difusión desde los streams de GCS
+        logger.info(f"Cargando modelo Diffusers para '{model_name}'...")
+        pipe = StableDiffusionPipeline.from_pretrained(io.BytesIO(model_bin_stream.read()))
+    else:
+        raise HTTPException(status_code=404, detail="No se encontró modelo compatible en el bucket.")
+    
+    return pipe
+
+def load_audiocraft_model_from_streams(model_data, model_name):
+    model_bin_stream = model_data.get("model_bin")
+    model_safetensors_stream = model_data.get("model_safetensors")
+
+    if model_bin_stream or model_safetensors_stream:
+        # Cargar el modelo AudioCraft desde los streams de GCS
+        logger.info(f"Cargando modelo Audiocraft para '{model_name}'...")
+        model = AudioLM.from_pretrained(io.BytesIO(model_bin_stream.read()))
+    else:
+        raise HTTPException(status_code=404, detail="No se encontró modelo compatible en el bucket.")
+    
+    return model
+
 @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.")
+
+        model_data = load_model_from_gcs(model_prefix, gcs_handler)
         
-        # Tareas basadas en texto
-        if request.pipeline_task in ["text-generation", "translation", "summarization"]:
-            pipe = pipeline(request.pipeline_task, model=model_stream, tokenizer=tokenizer_stream)
+        # Cargar los archivos de modelo y tokenizer directamente desde los streams
+        config_stream = model_data["config"]
+        tokenizer_stream = model_data["tokenizer"]
+
+        if request.pipeline_task == "text-generation":
+            # Usar el modelo HuggingFace normal si es una tarea de texto
+            model = load_model_from_streams(model_data, model_prefix)
+            tokenizer = AutoTokenizer.from_pretrained(io.BytesIO(tokenizer_stream.read()))
+            pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
             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.")
-            
+            # Usar el modelo Diffuser si es tarea de generación de imágenes
+            pipe = load_diffuser_model_from_streams(model_data, model_prefix)
+            result = pipe(request.input_text).images
+        elif request.pipeline_task == "audio-generation":
+            # Usar el modelo Audiocraft si es tarea de generación de audio
+            model = load_audiocraft_model_from_streams(model_data, model_prefix)
+            result = model.generate(request.input_text)
+        elif request.pipeline_task == "text-to-speech":
+            # TTS pipeline utilizando transformers
+            tts_pipe = tts_pipeline("text-to-speech", model=model, tokenizer=tokenizer)
+            audio_output = tts_pipe(request.input_text)[0]['audio']
+            # Se devuelve el archivo de audio
+            audio_path = "output.wav"
+            sf.write(audio_path, audio_output, 16000)  # Guardar el audio en un archivo
+            result = audio_path
+        elif request.pipeline_task == "text-to-audio":
+            # Usar audiocraft o modelo específico para text-to-audio
+            model = load_audiocraft_model_from_streams(model_data, model_prefix)
+            audio_output = model.generate(request.input_text)
+            # Guardar o procesar el audio de salida
+            audio_path = "output_audio.wav"
+            sf.write(audio_path, audio_output, 16000)
+            result = audio_path
+        else:
+            raise HTTPException(status_code=400, detail="Tarea no soportada.")
+        
+        logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
+        return {"response": result[0]}
+
     except HTTPException as e:
         logger.error(f"HTTPException: {e.detail}")
         raise e
@@ -212,29 +207,18 @@ async def predict(request: DownloadModelRequest):
         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"
+def download_model_in_background(model_name):
     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)
+        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+        logger.info(f"Iniciando descarga en segundo plano del modelo '{model_name}' a GCS...")
+        download_model_from_huggingface(model_name)
+        logger.info(f"Descarga del modelo '{model_name}' completada.")
     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.")
+        logger.error(f"Error al descargar el modelo '{model_name}' en segundo plano: {e}")
 
-# 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()
+    threading.Thread(target=download_model_in_background, args=("modelo_ejemplo",)).start()
 
 @app.on_event("startup")
 async def startup_event():
@@ -242,4 +226,3 @@ async def startup_event():
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
-