5

Dockerfile

@@ -1,28 +1,30 @@
-# Usar Python base
-FROM python:3.10-slim
-
-# Instalar dependencias del sistema
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    libgl1-mesa-glx \
-    libglib2.0-0 \
+# Usa una imagen base con soporte para Python y CUDA
+FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
+
+# Establece variables de entorno necesarias
+ENV TRANSFORMERS_CACHE="/app/.cache" \
+    HF_HOME="/app/.cache" \
+    PATH="/opt/conda/bin:$PATH"
+
+# Instala dependencias básicas
+RUN apt-get update && apt-get install -y \
+    python3 \
+    python3-pip \
+    git \
     && rm -rf /var/lib/apt/lists/*
 
-# Instalar dependencias de Python específicas para Marigold
-RUN pip install --no-cache-dir diffusers==0.16.1 transformers==4.30.2 huggingface_hub==0.16.4 accelerate==0.21.0 fastapi uvicorn torch torchvision pillow
+# Establece el directorio de trabajo
+WORKDIR /app
 
+# Copia el archivo requirements.txt al contenedor
+COPY requirements.txt /app/requirements.txt
 
-# Crear directorio de caché con permisos
-RUN mkdir -p /app/cache && chmod -R 777 /app/cache
-ENV HF_HOME=/app/cache
+# Instala las dependencias necesarias
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
 
-# Establecer directorio de trabajo
-WORKDIR /app
-
-# Copiar los archivos al contenedor
+# Copia los archivos de la aplicación al contenedor
 COPY . /app
 
-# Exponer el puerto donde se ejecutará el servidor
-EXPOSE 7860
-
-# Comando para iniciar la aplicación
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+# Comando para ejecutar la aplicación
+CMD ["python3", "app.py"]

app.py

@@ -1,33 +1,28 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI
 from diffusers import StableDiffusionDepth2ImgPipeline
-from PIL import Image
 import torch
 
 app = FastAPI()
 
-# Cargar el modelo Marigold
-model_id = "prs-eth/marigold-v1-0"
-pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
-    model_id, torch_dtype=torch.float16
-)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+# Inicializa el modelo Marigold
+@app.on_event("startup")
+async def load_model():
+    global pipe
+    pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
+        "prs-eth/marigold-depth-v1-0",
+        torch_dtype=torch.float16
+    )
+    pipe.to("cuda")  # Mueve el modelo a la GPU
+    print("Modelo cargado correctamente")
 
-@app.post("/generate-depth/")
-async def generate_depth(file: UploadFile = File(...)):
-    """
-    Generar un mapa de profundidad a partir de una imagen.
-    """
-    try:
-        # Leer la imagen
-        image = Image.open(file.file).convert("RGB")
 
-        # Generar el mapa de profundidad
-        depth_map = pipe(image).images[0]
+@app.get("/")
+async def root():
+    return {"message": "Marigold API funcionando correctamente"}
 
-        # Guardar el resultado en un archivo temporal
-        depth_map.save("/app/cache/depth_map.png")
 
-        # Retornar un mensaje de éxito
-        return {"message": "Mapa de profundidad generado con éxito.", "output": "/app/cache/depth_map.png"}
-    except Exception as e:
-        return {"error": str(e)}
+@app.post("/depthmap/")
+async def generate_depth_map(input_image: str):
+    # Carga la imagen y genera el mapa de profundidad
+    image = pipe(input_image).images[0]
+    return {"depth_map": image}

requirements.txt

@@ -1,7 +1,10 @@
 fastapi
 uvicorn[standard]
-torch
-torchvision
-diffusers
-transformers
 pillow
+accelerate>=0.22.0
+diffusers>=0.25.0
+matplotlib
+scipy
+torch==2.4.0
+torchvision==0.19.0
+transformers>=4.32.1