From 16617340acb2155c6c3e9b42013ce2daf0ec4945 Mon Sep 17 00:00:00 2001
From: Banbury <banbury20@gmail.com>
Date: Wed, 4 Sep 2024 21:38:02 +0200
Subject: [PATCH 1/2] Added fastapi server and Docker files.

---
 Dockerfile         | 16 +++++++++++
 docker-compose.yml | 12 ++++++++
 env.server         |  7 +++++
 requirements.txt   |  3 +-
 server.py          | 72 ++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 109 insertions(+), 1 deletion(-)
 create mode 100644 Dockerfile
 create mode 100644 docker-compose.yml
 create mode 100644 env.server
 create mode 100644 server.py

diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..b7c83f3
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,16 @@
+FROM python:slim
+
+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')

From 8e9de478f047b16f30c97e70ed5bfed9b19852ba Mon Sep 17 00:00:00 2001
From: Banbury <banbury20@gmail.com>
Date: Tue, 17 Sep 2024 10:52:36 +0200
Subject: [PATCH 2/2] Rebased to master. Dwongraded Debian to Bullseye.

---
 Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index b7c83f3..ddaf27f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:slim
+FROM python:slim-bullseye
 
 WORKDIR /app
 RUN apt-get update && \