Update app.py

app.py

@@ -11,8 +11,8 @@ from google.auth import exceptions
 from transformers import pipeline
 from dotenv import load_dotenv
 import uvicorn
-import tempfile
 
+# Configuración de carga de variables de entorno
 load_dotenv()
 
 API_KEY = os.getenv("API_KEY")
@@ -20,9 +20,11 @@ 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)
@@ -63,7 +65,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")
+        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)
@@ -74,6 +76,7 @@ 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)
@@ -87,6 +90,7 @@ 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)
@@ -117,28 +121,125 @@ async def predict(request: DownloadModelRequest):
             logger.info(f"Modelos no encontrados en GCS, descargando '{model_prefix}' desde Hugging Face...")
             download_model_from_huggingface(model_prefix)
         
-        model_files_streams = {}
-        with tempfile.TemporaryDirectory() as temp_dir:
-            for file in model_files:
-                if gcs_handler.file_exists(f"{model_prefix}/{file}"):
-                    file_path = os.path.join(temp_dir, file)
-                    with open(file_path, "wb") as f:
-                        gcs_handler.download_file(f"{model_prefix}/{file}").readinto(f)
-                    model_files_streams[file] = file_path
-            
-            if not all(key in model_files_streams for key in ["config.json", "tokenizer.json", "pytorch_model.bin"]):
-                logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
-                raise HTTPException(status_code=500, detail="Required model files missing.")
+        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.")
             
-            if request.pipeline_task in ["text-generation", "translation", "summarization"]:
-                pipe = pipeline(request.pipeline_task, model=model_files_streams["pytorch_model.bin"], tokenizer=model_files_streams["tokenizer.json"])
-                result = pipe(request.input_text)
-                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
+    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 en la predicción: {e}")
-        raise HTTPException(status_code=500, detail=f"Error en la predicción: {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)
+