diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ddaf27f --- /dev/null +++ b/Dockerfile @@ -0,0 +1,16 @@ +FROM python:slim-bullseye + +WORKDIR /app +RUN apt-get update && \ + apt-get upgrade -y && \ + apt-get install -y git build-essential python3-setuptools +RUN git clone https://github.com/Stability-AI/stable-fast-3d.git +WORKDIR /app/stable-fast-3d +COPY server.py . +COPY env.server .env +COPY requirements.txt . +RUN mkdir model +RUN rm __init__.py +RUN python -m pip install torch torchvision torchaudio setuptools==69.5.1 wheel +RUN python -m pip install -r requirements.txt +CMD ["fastapi", "run", "/app/stable-fast-3d/server.py", "--port", "8000"] diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..9b9db87 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,12 @@ +version: '3.4' + +services: + stablefast3d: + image: stablefast3d + ports: + - 8100:8000 + build: + context: . + dockerfile: ./Dockerfile + volumes: + - ./model:/app/stable-fast-3d/model diff --git a/env.server b/env.server new file mode 100644 index 0000000..67631ce --- /dev/null +++ b/env.server @@ -0,0 +1,7 @@ +DEVICE=cpu +PRETRAINED_MODEL=model +FOREGROUND_RATIO=0.85 +TEXTURE_RESOLUTION=1024 +REMESSH_OPTION=none +TARGET_VERTEX_COUNT=-1 +BATCH_SIZE=1 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 47be34c..8f46ae1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,6 +9,7 @@ huggingface-hub==0.23.4 rembg[gpu]==2.0.57; sys_platform != 'darwin' rembg==2.0.57; sys_platform == 'darwin' git+https://github.com/vork/PyNanoInstantMeshes.git@v0.0.3 -gpytoolbox==0.2.0 +gpytoolbox==0.3.2 +fastapi[standard]>=0.112.2 ./texture_baker/ ./uv_unwrapper/ diff --git a/server.py b/server.py new file mode 100644 index 0000000..ff7696a --- /dev/null +++ b/server.py @@ -0,0 +1,72 @@ +import base64 +import os +import io +from io import BytesIO +from contextlib import nullcontext +from dotenv import load_dotenv +from fastapi import FastAPI, UploadFile +from fastapi.responses import Response +from PIL import Image +import rembg +import torch +from sf3d.system import SF3D +from sf3d.utils import get_device, remove_background, resize_foreground + +load_dotenv() + +app = FastAPI() + +model = SF3D.from_pretrained( + os.getenv("PRETRAINED_MODEL"), + config_name="config.yaml", + weight_name="model.safetensors", +) +model.to(os.getenv("DEVICE")) +model.eval() + +rembg_session = rembg.new_session() + +@app.post("/generate", + responses = { 200: { "content": { "model/gltf-binary": {} } } }, + response_class=Response +) +async def generate(file: UploadFile): + # load the image using Pillow + image = Image.open(file.file) + image = remove_bg(image) + mesh = generate_model(image) + + # return the image as a binary stream with a suitable content-disposition header for download + return Response( + content=mesh, + media_type="model/gltf-binary", + headers={"Content-Disposition": "attachment; filename=mesh.glb"}, + ) + +def remove_bg(img: Image) -> Image: + img = remove_background(img, rembg_session) + img = resize_foreground(img, float(os.getenv("FOREGROUND_RATIO"))) + return img + +def generate_model(image: Image): + device = os.getenv("DEVICE") + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + with torch.no_grad(): + with torch.autocast( + device_type=device, dtype=torch.float16 + ) if "cuda" in device else nullcontext(): + mesh, glob_dict = model.run_image( + image, + bake_resolution=int(os.getenv("TEXTURE_RESOLUTION")), + remesh=os.getenv("REMESSH_OPTION"), + vertex_count=int(os.getenv("TARGET_VERTEX_COUNT")), + ) + if torch.cuda.is_available(): + print("Peak Memory:", torch.cuda.max_memory_allocated() / 1024 / 1024, "MB") + elif torch.backends.mps.is_available(): + print( + "Peak Memory:", torch.mps.driver_allocated_memory() / 1024 / 1024, "MB" + ) + + return mesh.export(include_normals=True, file_type='glb')