crevelop commited on
Commit
3724212
·
unverified ·
1 Parent(s): 12b2ede

add api endpoint

Browse files
Files changed (1) hide show
  1. app.py +65 -1
app.py CHANGED
@@ -16,7 +16,10 @@ from trellis.representations import Gaussian, MeshExtractResult
16
  from trellis.utils import render_utils, postprocessing_utils
17
 
18
  import logging
19
-
 
 
 
20
  # Configure logging
21
  logging.basicConfig(
22
  level=logging.INFO,
@@ -264,6 +267,67 @@ with gr.Blocks() as demo:
264
  )
265
 
266
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
267
  # Launch the Gradio app
268
  if __name__ == "__main__":
269
  pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
 
16
  from trellis.utils import render_utils, postprocessing_utils
17
 
18
  import logging
19
+ from fastapi import FastAPI, File, UploadFile
20
+ from fastapi.responses import JSONResponse, FileResponse
21
+ import io
22
+ import base64
23
  # Configure logging
24
  logging.basicConfig(
25
  level=logging.INFO,
 
267
  )
268
 
269
 
270
+ # Access FastAPI app from Gradio
271
+ app = gr.routes.App.get_app()
272
+
273
+ @app.post("/api/image_to_glb")
274
+ async def process_image_endpoint(
275
+ image: UploadFile = File(...),
276
+ seed: int = 0,
277
+ randomize_seed: bool = True,
278
+ ss_guidance_strength: float = 7.5,
279
+ ss_sampling_steps: int = 12,
280
+ slat_guidance_strength: float = 3.0,
281
+ slat_sampling_steps: int = 12
282
+ ):
283
+ """
284
+ API Endpoint to process an image and return a GLB file.
285
+
286
+ Args:
287
+ image (UploadFile): The image file.
288
+ seed (int): Seed for generation.
289
+ randomize_seed (bool): Whether to randomize the seed.
290
+ ss_guidance_strength (float): Guidance strength for stage 1.
291
+ ss_sampling_steps (int): Sampling steps for stage 1.
292
+ slat_guidance_strength (float): Guidance strength for stage 2.
293
+ slat_sampling_steps (int): Sampling steps for stage 2.
294
+
295
+ Returns:
296
+ FileResponse: The generated GLB file as a downloadable attachment.
297
+ """
298
+ try:
299
+ # Read and preprocess the image
300
+ contents = await image.read()
301
+ pil_image = Image.open(io.BytesIO(contents)).convert("RGBA")
302
+ trial_id, processed_image = preprocess_image(pil_image)
303
+
304
+ # Generate 3D asset
305
+ state, video_path = image_to_3d(
306
+ trial_id, seed, randomize_seed,
307
+ ss_guidance_strength, ss_sampling_steps,
308
+ slat_guidance_strength, slat_sampling_steps
309
+ )
310
+
311
+ # Extract GLB
312
+ glb_path, _ = extract_glb(state, mesh_simplify=0.95, texture_size=1024) # You can parametrize these
313
+
314
+ # Ensure the GLB file exists
315
+ if not os.path.exists(glb_path):
316
+ logger.error(f"GLB file not found at path: {glb_path}")
317
+ return JSONResponse(status_code=500, content={"error": "GLB file generation failed."})
318
+
319
+ # Return the GLB file as a downloadable response
320
+ return FileResponse(
321
+ path=glb_path,
322
+ media_type='model/gltf-binary',
323
+ filename=f"{trial_id}.glb"
324
+ )
325
+
326
+ except Exception as e:
327
+ logger.error(f"Error in API endpoint: {e}")
328
+ return JSONResponse(status_code=500, content={"error": str(e)})
329
+
330
+
331
  # Launch the Gradio app
332
  if __name__ == "__main__":
333
  pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")