Update app.py

app.py

@@ -1,12 +1,15 @@
-# app.py (simplified)
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+import os
+import tempfile
+import numpy as np
 import torch
+from flask import Flask, request, jsonify, send_file
+from flask_cors import CORS
 from point_e.diffusion.sampler import PointCloudSampler
 from point_e.models.configs import MODEL_CONFIGS, model_from_config
-from point_e.util.plotting import save_point_cloud
-import tempfile
-import os
+import open3d as o3d
+
+# Fix matplotlib temp issue
+os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
 
 app = Flask(__name__)
 CORS(app)
@@ -18,26 +21,32 @@ base_model = model_from_config(MODEL_CONFIGS["base40M-textvec"], device)
 base_model.load_state_dict(torch.load("base40M-textvec.pt", map_location=device))
 base_model.eval()
 
-print("Loading sampler...")
+print("Creating sampler...")
 sampler = PointCloudSampler(
     device=device,
     models=[base_model],
     diffusion=None,
     num_points=1024,
     aux_channels=[],
-    guidance_scale=1.0,
+    guidance_scale=3.0,
 )
 
+def save_point_cloud_open3d(pc_xyz: np.ndarray, out_path: str):
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(pc_xyz)
+    o3d.io.write_point_cloud(out_path, pcd)
+
 @app.route("/generate", methods=["POST"])
 def generate():
-    prompt = request.json.get("prompt", "a red apple")
-    samples = sampler.sample_batch(prompt=[prompt])
+    prompt = request.json.get("prompt", "a toy robot")
+    samples = sampler.sample_batch(batch_size=1, model_kwargs={"texts": [prompt]})
     pc = sampler.output_to_point_clouds(samples)[0]
 
     tmp_dir = tempfile.mkdtemp()
     output_path = os.path.join(tmp_dir, "model.ply")
-    save_point_cloud(pc, output_path)
-    return jsonify({"model_path": output_path})
+    save_point_cloud_open3d(pc['coords'], output_path)
+
+    return send_file(output_path, as_attachment=True)
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)