davidvgilmore commited on
Commit
5cf2896
·
verified ·
1 Parent(s): 4e1f061

Upload api.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. api.py +83 -0
api.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Open Source Model Licensed under the Apache License Version 2.0
2
+ # and Other Licenses of the Third-Party Components therein:
3
+ # The below Model in this distribution may have been modified by THL A29 Limited
4
+ # ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
5
+
6
+ import os
7
+ import torch
8
+ from PIL import Image
9
+ from fastapi import FastAPI, File, UploadFile, HTTPException
10
+ from fastapi.responses import FileResponse
11
+ from pydantic import BaseModel
12
+ import io
13
+ import tempfile
14
+
15
+ from hy3dgen.rembg import BackgroundRemover
16
+ from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline, FaceReducer, FloaterRemover, DegenerateFaceRemover
17
+ from hy3dgen.text2image import HunyuanDiTPipeline
18
+
19
+ app = FastAPI(title="Hunyuan3D-2 API")
20
+
21
+ # Initialize models at startup
22
+ rembg = BackgroundRemover()
23
+ model_path = 'tencent/Hunyuan3D-2' # Use Hugging Face model path
24
+ t2i = HunyuanDiTPipeline('Tencent-Hunyuan--HunyuanDiT-v1.1-Diffusers-Distilled')
25
+ i23d = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(model_path)
26
+
27
+ class TextToImageRequest(BaseModel):
28
+ prompt: str
29
+
30
+ @app.post("/image-to-3d/")
31
+ async def image_to_3d(file: UploadFile = File(...)):
32
+ """Convert an image to 3D model"""
33
+ try:
34
+ # Read and process image
35
+ content = await file.read()
36
+ image = Image.open(io.BytesIO(content))
37
+ image = image.resize((1024, 1024))
38
+
39
+ if image.mode == 'RGB':
40
+ image = rembg(image)
41
+
42
+ # Generate mesh
43
+ mesh = i23d(image=image, num_inference_steps=30, mc_algo='mc',
44
+ generator=torch.manual_seed(2025))[0]
45
+ mesh = FloaterRemover()(mesh)
46
+ mesh = DegenerateFaceRemover()(mesh)
47
+ mesh = FaceReducer()(mesh)
48
+
49
+ # Save mesh to temporary file
50
+ with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
51
+ mesh.export(tmp.name)
52
+ return FileResponse(tmp.name, media_type='model/gltf-binary',
53
+ filename='model.glb')
54
+
55
+ except Exception as e:
56
+ raise HTTPException(status_code=500, detail=str(e))
57
+
58
+ @app.post("/text-to-3d/")
59
+ async def text_to_3d(request: TextToImageRequest):
60
+ """Convert text prompt to 3D model"""
61
+ try:
62
+ # Generate image from text
63
+ image = t2i(request.prompt)
64
+ image = rembg(image)
65
+
66
+ # Generate mesh
67
+ mesh = i23d(image, num_inference_steps=30, mc_algo='mc')[0]
68
+ mesh = FloaterRemover()(mesh)
69
+ mesh = DegenerateFaceRemover()(mesh)
70
+ mesh = FaceReducer()(mesh)
71
+
72
+ # Save mesh to temporary file
73
+ with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
74
+ mesh.export(tmp.name)
75
+ return FileResponse(tmp.name, media_type='model/gltf-binary',
76
+ filename='model.glb')
77
+
78
+ except Exception as e:
79
+ raise HTTPException(status_code=500, detail=str(e))
80
+
81
+ if __name__ == "__main__":
82
+ import uvicorn
83
+ uvicorn.run(app, host="0.0.0.0", port=8000)