-
Notifications
You must be signed in to change notification settings - Fork 155
Commit
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
FROM python:slim | ||
Check failure on line 1 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerMissing User Instruction
Raw output
Check notice on line 1 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerHealthcheck Instruction Missing
Raw output
|
||
|
||
WORKDIR /app | ||
RUN apt-get update && \ | ||
Check warning on line 4 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerApt Get Install Pin Version Not Defined
Raw output
Check warning on line 4 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerApt Get Install Pin Version Not Defined
Raw output
Check warning on line 4 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerApt Get Install Pin Version Not Defined
Raw output
Check notice on line 4 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerAPT-GET Not Avoiding Additional Packages
Raw output
Check notice on line 4 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerApt Get Install Lists Were Not Deleted
Raw output
|
||
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 | ||
Check notice on line 12 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerMultiple RUN, ADD, COPY, Instructions Listed
Raw output
|
||
RUN rm __init__.py | ||
RUN python -m pip install torch torchvision torchaudio setuptools==69.5.1 wheel | ||
Check warning on line 14 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerPip install keeping cached packages
Raw output
Check warning on line 14 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerUnpinned Package Version in Pip Install
Raw output
|
||
RUN python -m pip install -r requirements.txt | ||
Check warning on line 15 in Dockerfile Wiz Inc. (a28a8b7b4c) / Wiz IaC ScannerPip install keeping cached packages
Raw output
|
||
CMD ["fastapi", "run", "/app/stable-fast-3d/server.py", "--port", "8000"] |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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') |