MarigoldDepthPipeline

Dockerfile

@@ -1,38 +1,30 @@
-# Usa una imagen base con soporte para Python y CUDA
+# Imagen base para CUDA y PyTorch
 FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
 
-# Instala dependencias básicas
+# Configuración de entornos
+ENV TRANSFORMERS_CACHE="/app/.cache" \
+    HF_HOME="/app/.cache"
+
+# Instalación de herramientas básicas
 RUN apt-get update && apt-get install -y \
     python3 \
     python3-pip \
     git \
     && rm -rf /var/lib/apt/lists/*
 
-# Establece el directorio de trabajo
-WORKDIR /app
-
-# Clona el repositorio de Marigold
-RUN git clone https://github.com/prs-eth/Marigold.git /app/Marigold
-
-# Instala las dependencias del repositorio de Marigold
+# Instala dependencias de Python
 RUN pip install --upgrade pip
-RUN pip install -r /app/Marigold/requirements.txt
+RUN pip install torch torchvision diffusers fastapi uvicorn pillow prs-eth-marigold
 
-# Copia el archivo requirements.txt para dependencias adicionales (API, etc.)
-COPY requirements.txt /app/requirements.txt
-
-# Instala las dependencias necesarias para la API
-RUN pip install -r /app/requirements.txt
-
-# Copia los archivos necesarios para la API
-COPY app.py /app/app.py
-COPY entrypoint.sh /app/entrypoint.sh
+# Crea directorios necesarios
+WORKDIR /app
+RUN mkdir -p /app/.cache && chmod -R 777 /app/.cache
 
-# Da permisos de ejecución al archivo entrypoint
-RUN chmod +x /app/entrypoint.sh
+# Copia el código de la aplicación
+COPY . /app
 
-# Exponer el puerto para la API
+# Expone el puerto de FastAPI
 EXPOSE 8000
 
-# Comando para inicializar la app
-CMD ["/app/entrypoint.sh"]
+# Comando de inicio
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app.py

@@ -1,36 +1,31 @@
-from fastapi import FastAPI, HTTPException
-import subprocess
-import os
+from fastapi import FastAPI, UploadFile, File
+from PIL import Image
+from prs_eth_marigold import MarigoldDepthPipeline
 
 app = FastAPI()
 
-@app.post("/run_inference/")
-def run_inference(input_dir: str = "Marigold/input/in-the-wild_example", output_dir: str = "Marigold/output"):
+# Carga el modelo Marigold al inicio
+pipe = MarigoldDepthPipeline.from_pretrained(
+    "prs-eth/marigold-depth-v1-0", torch_dtype="auto"
+).to("cuda")
+
+@app.post("/generate-depth/")
+async def generate_depth(file: UploadFile = File(...)):
     """
-    Llama al script de inferencia de Marigold con los parámetros proporcionados.
+    Genera un mapa de profundidad a partir de una imagen.
     """
     try:
-        # Verifica que los directorios existan
-        if not os.path.exists(input_dir):
-            raise HTTPException(status_code=400, detail="Directorio de entrada no encontrado.")
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
+        image = Image.open(file.file)
 
-        # Llama al script de inferencia
-        command = [
-            "python", "Marigold/run.py",
-            "--checkpoint", "prs-eth/marigold-v1-0",
-            "--denoise_steps", "50",
-            "--ensemble_size", "10",
-            "--input_rgb_dir", input_dir,
-            "--output_dir", output_dir
-        ]
-        result = subprocess.run(command, capture_output=True, text=True)
+        result = pipe(image).images[0]
 
-        # Verifica si hubo errores durante la ejecución
-        if result.returncode != 0:
-            raise HTTPException(status_code=500, detail=f"Error en la inferencia: {result.stderr}")
+        output_path = "/app/output_depth_map.png"
+        result.save(output_path)
 
-        return {"message": "Inferencia completada", "output_dir": output_dir}
+        return {"message": "Mapa de profundidad generado", "output_path": output_path}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return {"error": str(e)}
+    
+@app.get("/")
+async def root():
+    return {"message": "API de generación de mapas de profundidad con Marigold"}

requirements.txt

@@ -1,2 +1,9 @@
 fastapi>=0.95.0
 uvicorn[standard]
+accelerate>=0.22.0
+diffusers>=0.25.0
+matplotlib
+scipy
+torch==2.0.1
+torchvision==0.15.2
+transformers>=4.32.1