add api routes

app.py

@@ -138,6 +138,7 @@ def image_to_3d(
         dict: The information of the generated 3D model.
         str: The path to the video of the 3D model.
     """
+    
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     if not is_multiimage:
         outputs = pipeline.run(
@@ -180,6 +181,43 @@ def image_to_3d(
     return state, video_path
 
 
+@spaces.GPU
+def image_to_3d2(
+    image: Image.Image,
+    seed: int,
+    ss_guidance_strength: float,
+    ss_sampling_steps: int,
+    slat_guidance_strength: float,
+    slat_sampling_steps: int,
+) -> Tuple[dict, str]:
+    
+    outputs = pipeline.run(
+        image,
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+    )
+
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
+        video_path = temp_video.name
+        imageio.mimsave(video_path, video, fps=15)
+    
+    torch.cuda.empty_cache()
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+    return state, video_path
+
+import random
 @spaces.GPU(duration=90)
 def extract_glb(
     state: dict,
@@ -198,10 +236,11 @@ def extract_glb(
     Returns:
         str: The path to the extracted GLB file.
     """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    user_dir = TMP_DIR
     gs, mesh = unpack_state(state)
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(user_dir, 'sample.glb')
+    glb_path = os.path.join(user_dir, f"test_{random.random()}.glb")
+
     glb.export(glb_path)
     torch.cuda.empty_cache()
     return glb_path, glb_path
@@ -254,6 +293,57 @@ def split_image(image: Image.Image) -> List[Image.Image]:
         images.append(Image.fromarray(image[:, s:e+1]))
     return [preprocess_image(image) for image in images]
 
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse, FileResponse
+import tempfile
+
+app = FastAPI()
+
+@app.post("/generatee")
+async def generate_3d(image: UploadFile = File(...)):
+    if not image:
+        raise HTTPException(status_code=400, detail="No image provided")
+
+    try:
+        image_data = Image.open(image.file)
+        image_data = image_data.convert("RGBA")
+        # image_data = torch.tensor(np.array(image_data)).to('cuda')
+
+        seed = 42
+        ss_guidance_strength = 7.5
+        ss_sampling_steps = 12
+        slat_guidance_strength = 3.0
+        slat_sampling_steps = 12
+
+        print("STARTING TO GENERATE VIDEO ...")
+        state, _ = image_to_3d2(
+            image_data,
+            seed=seed,
+            ss_guidance_strength=ss_guidance_strength,
+            ss_sampling_steps=ss_sampling_steps,
+            slat_guidance_strength=slat_guidance_strength,
+            slat_sampling_steps=slat_sampling_steps,
+        )
+        mesh_simplify = 0.95
+        texture_size = 1024
+        print("STARTING TO GENERATE GLB FILE ...")
+        glb_path, _ = extract_glb(
+            state=state,
+            mesh_simplify=mesh_simplify,
+            texture_size=texture_size,
+            req=None  # Assuming req is not needed here
+        )
+        print("DONE")
+        return FileResponse(glb_path, media_type='application/octet-stream', filename='model.glb')
+
+    except Exception as e:
+        print("ERROR IN GENERATING 3D FILE : " , e)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/test")
+async def test_():
+    return JSONResponse(content={"TEST": "DONE!"})
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
@@ -405,10 +495,17 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
 # Launch the Gradio app
 if __name__ == "__main__":
+    import uvicorn
+    from gradio import mount_gradio_app
+    
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
-    try:
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
-    except:
-        pass
-    demo.launch()
+    
+    app = mount_gradio_app(app, demo, path="/")
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+
+    # try:
+    #     pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
+    # except:
+    #     pass
+    # demo.launch()