Update app.py

app.py

@@ -1,3 +1,7 @@
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from typing import Optional
 import logging
 import os
 import boto3
@@ -17,6 +21,9 @@ from functools import partial
 import io
 from io import BytesIO
 
+app = FastAPI()
+
+torch.cuda.empty_cache()
 
 subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
 
@@ -165,7 +172,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
             image = fill_background(image)
     return image
 
-@spaces.GPU
+# @spaces.GPU
 def generate(image, mc_resolution, formats=["obj", "glb"]):
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
@@ -180,73 +187,32 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     
     return mesh_path_obj.name, mesh_path_glb.name
 
-def run_example(image, seed, do_remove_background, foreground_ratio, mc_resolution, text_prompt=None):
-    image_pil = generate_image_from_text(encoded_image=image, seed=seed, pos_prompt=text_prompt)
+@app.post("/process_image/")
+async def process_image(
+    file: UploadFile = File(...),
+    seed: int = Form(...),
+    use_image: bool = Form(...),
+    do_remove_background: bool = Form(...),
+    foreground_ratio: float = Form(...),
+    mc_resolution: int = Form(...),
+    text_prompt: Optional[str] = Form(None)
+):
+    image_bytes = await file.read()
+    input_image = Image.open(BytesIO(image_bytes))
+
+    if use_image:
+        image_pil = generate_image_from_text(encoded_image=input_image, seed=seed, pos_prompt=text_prompt)
+    else:
+        image_pil = input_image
+
     preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
-    mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
-    return preprocessed, mesh_name_obj, mesh_name_glb
+    mesh_name_obj, mesh_name_glb = generate(preprocessed, mc_resolution)
 
-    
-with gr.Blocks() as demo:
-    gr.Markdown(HEADER)
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                
-                input_image = gr.Image(
-                    label="Generated Image",
-                    image_mode="RGBA",
-                    sources="upload",
-                    type="pil",
-                    elem_id="content_image" 
-                )
-                text_prompt = gr.Textbox(
-                    label="Text Prompt",
-                    placeholder="Enter Positive Prompt"
-                )
-                seed = gr.Textbox(label="Random Seed", value=0)
-                processed_image = gr.Image(label="Processed Image", interactive=False, visible=False)
-            with gr.Row():
-                with gr.Group():
-                    do_remove_background = gr.Checkbox(
-                        label="Remove Background", value=True
-                    )
-                    foreground_ratio = gr.Slider(
-                        label="Foreground Ratio",
-                        minimum=0.5,
-                        maximum=1.0,
-                        value=0.85,
-                        step=0.05,
-                    )
-                    mc_resolution = gr.Slider(
-                        label="Marching Cubes Resolution",
-                        minimum=32,
-                        maximum=320,
-                        value=256,
-                        step=32
-                     )
-            with gr.Row():
-                submit = gr.Button("Generate", elem_id="generate", variant="primary")
-        with gr.Column():
-            with gr.Tab("OBJ"):
-                output_model_obj = gr.Model3D(
-                    label="Output Model (OBJ Format)",
-                    interactive=False,
-                )
-                gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
-            with gr.Tab("GLB"):
-                output_model_glb = gr.Model3D(
-                    label="Output Model (GLB Format)",
-                    interactive=False,
-                )
-                gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
-                
-    submit.click(fn=check_input_image, inputs=[input_image]).success(
-        fn=run_example,
-        inputs=[input_image, seed, do_remove_background, foreground_ratio, mc_resolution, text_prompt],
-        outputs=[processed_image, output_model_obj, output_model_glb],
-        # outputs=[output_model_obj, output_model_glb],
-    )
-
-demo.queue(max_size=10)
-demo.launch()
+    return {
+        "obj_path": mesh_name_obj,
+        "glb_path": mesh_name_glb
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)