Update app.py

app.py

@@ -1,11 +1,9 @@
-# app.py
+# app.py (simplified)
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 import torch
-from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
-from point_e.models.download import load_checkpoint
-from point_e.models.text_to_image import create_model_and_diffusion as create_t2i
-from point_e.models.image_to_3d import create_model_and_diffusion as create_i2p
+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
@@ -13,54 +11,33 @@ import os
 app = Flask(__name__)
 CORS(app)
 
-# Load models once on startup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-print(f"Using device: {device}")
 
-print("Loading text-to-image model...")
-t2i_model, t2i_diffusion = create_t2i(DIFFUSION_CONFIGS['base40'], device)
-t2i_model.load_state_dict(load_checkpoint('base40', device))
-t2i_model.eval()
-
-print("Loading image-to-3d model...")
-i2p_model, i2p_diffusion = create_i2p(DIFFUSION_CONFIGS['base40'], device)
-i2p_model.load_state_dict(load_checkpoint('base40-img2pc', device))
-i2p_model.eval()
-
-@app.route('/generate', methods=['POST'])
+print("Loading base model...")
+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...")
+sampler = PointCloudSampler(
+    device=device,
+    models=[base_model],
+    diffusion=None,
+    num_points=1024,
+    aux_channels=[],
+    guidance_scale=1.0,
+)
+
+@app.route("/generate", methods=["POST"])
 def generate():
-    data = request.json
-    prompt = data.get('prompt', 'a red car')
-
-    # Generate image from text
-    print("Generating image from text...")
-    samples = t2i_diffusion.sample_loop(
-        model=t2i_model,
-        shape=(1, 3, 64, 64),
-        device=device,
-        clip_denoised=True,
-        model_kwargs={"text": [prompt]},
-        progress=True
-    )
-    image = t2i_model.output_to_image(samples[0])
+    prompt = request.json.get("prompt", "a red apple")
+    samples = sampler.sample_batch(prompt=[prompt])
+    pc = sampler.output_to_point_clouds(samples)[0]
 
-    # Generate point cloud from image
-    print("Generating point cloud...")
-    pc = i2p_diffusion.sample_loop(
-        model=i2p_model,
-        shape=(1, 3, 64, 64),
-        device=device,
-        clip_denoised=True,
-        model_kwargs={"images": [image]},
-        progress=True
-    )
-
-    # Save to temporary file
     tmp_dir = tempfile.mkdtemp()
-    output_path = os.path.join(tmp_dir, 'output.ply')
-    save_point_cloud(pc[0], output_path)
-
-    return jsonify({'model_path': output_path})
+    output_path = os.path.join(tmp_dir, "model.ply")
+    save_point_cloud(pc, output_path)
+    return jsonify({"model_path": output_path})
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)