Spaces:

mac9087
/

magical-box

Running

App Files Files Community

mac9087 commited on 23 days ago

Commit

ca7e93f

verified ·

1 Parent(s): 97be4f4

Update app.py

Browse files

Files changed (1) hide show

app.py +87 -119

app.py CHANGED Viewed

@@ -1,136 +1,104 @@
-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
-import open3d as o3d
-import logging
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-# Fix matplotlib and clip cache path issues
-os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
-os.environ['HOME'] = '/tmp'
-# Set cache directories explicitly
-cache_dir = os.environ.get('POINT_E_CACHE_DIR', '/tmp/point_e_model_cache')
-clip_cache_dir = os.environ.get('CLIP_MODEL_DIR', '/tmp/clip_models')
-# Ensure cache directories exist and are writable
-for directory in [cache_dir, clip_cache_dir, '/tmp/point_e_models']:
-    if not os.path.exists(directory):
-        os.makedirs(directory, exist_ok=True)
-    # Make sure permissions are correct
-    os.chmod(directory, 0o777)
 app = Flask(__name__)
 CORS(app)
-# Determine if CUDA is available (but prepare for CPU operation)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-logger.info(f"Using device: {device}")
-# Set up model directory for downloads
-model_path = "/tmp/point_e_models/base40M-textvec.pt"
-if not os.path.exists(model_path):
-    logger.info("Model weights not found. Downloading...")
-    # Use torch.hub to download the model from original repo
-    # This is a placeholder. In reality, you'd need to download the weights from the official source
-    # or include them in your repository
-    try:
-        # For demonstration purposes, we'll assume the model is included with the repo
-        # In a real setup, you'd need to download it or include it
-        pass
-    except Exception as e:
-        logger.error(f"Failed to download model: {e}")
-        raise
-# Load the model
-logger.info("Loading base model...")
-try:
-    base_model = model_from_config(
-        MODEL_CONFIGS["base40M-textvec"],
-        device=device,
-        cache_dir=cache_dir  # Pass cache directory explicitly
-    )
-    base_model.load_state_dict(torch.load(model_path, map_location=device))
-    base_model.eval()
-    logger.info("Base model loaded successfully")
-except Exception as e:
-    logger.error(f"Failed to load model: {e}")
-    raise
-# Create sampler with memory-efficient settings for 2vCPU/18GB RAM
-logger.info("Creating sampler...")
-sampler = PointCloudSampler(
-    device=device,
-    models=[base_model],
-    diffusion=None,
-    num_points=1024,  # Lower point count for memory efficiency
-    aux_channels=[],
-    guidance_scale=3.0,
-)
-def save_point_cloud_open3d(pc_xyz: np.ndarray, out_path: str):
-    """Save a point cloud to a PLY file using Open3D."""
-    pcd = o3d.geometry.PointCloud()
-    pcd.points = o3d.utility.Vector3dVector(pc_xyz)
-    o3d.io.write_point_cloud(out_path, pcd)
-    logger.info(f"Point cloud saved to {out_path}")
-@app.route("/", methods=["GET"])
-def index():
-    """Provide a simple interface for the API."""
-    return """
-    <html>
-        <head><title>Point-E Text-to-3D API</title></head>
-        <body>
-            <h1>Point-E Text-to-3D Generator</h1>
-            <form action="/generate" method="post">
-                <label for="prompt">Enter a description:</label>
-                <input type="text" id="prompt" name="prompt" value="a toy robot">
-                <input type="submit" value="Generate 3D Model">
-            </form>
-            <p>Or use the API endpoint at POST /generate with JSON {"prompt": "your description"}</p>
-        </body>
-    </html>
-    """
-@app.route("/generate", methods=["POST"])
-def generate():
-    """Generate a 3D point cloud from a text description."""
-    if request.is_json:
-        prompt = request.json.get("prompt", "a toy robot")
-    else:
-        prompt = request.form.get("prompt", "a toy robot")
-    logger.info(f"Generating point cloud for prompt: '{prompt}'")
     try:
-        # Generate with torch.no_grad() to save memory
-        with torch.no_grad():
-            samples = sampler.sample_batch(batch_size=1, model_kwargs={"texts": [prompt]})
-            pc = sampler.output_to_point_clouds(samples)[0]
-        # Save the point cloud
-        tmp_dir = tempfile.mkdtemp()
-        output_path = os.path.join(tmp_dir, f"{prompt.replace(' ', '_')}.ply")
-        save_point_cloud_open3d(pc['coords'], output_path)
-        return send_file(output_path, as_attachment=True, download_name=f"{prompt.replace(' ', '_')}.ply")
     except Exception as e:
-        logger.error(f"Error generating point cloud: {e}")
         return jsonify({"error": str(e)}), 500
-@app.route("/health", methods=["GET"])
-def health_check():
-    """Simple health check endpoint."""
-    return jsonify({"status": "healthy"}), 200
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)

 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
+import torch
+import numpy as np
+import trimesh
+import os
+from io import BytesIO
+import base64
+from PIL import Image
+import uuid
+import time
+# Import Shap-E
+from shap_e.diffusion.sample import sample_latents
+from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
+from shap_e.models.download import load_model, load_config
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_mesh
 app = Flask(__name__)
 CORS(app)
+# Create output directory if it doesn't exist
+os.makedirs("outputs", exist_ok=True)
+# Load models only once at startup
+print("Loading models...")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using device: {device}")
+xm = load_model('transmitter', device=device)
+model = load_model('text300M', device=device)
+diffusion = diffusion_from_config(load_config('diffusion'))
+print("Models loaded successfully!")
+@app.route('/generate', methods=['POST'])
+def generate_3d():
     try:
+        # Get the prompt from the request
+        data = request.json
+        if not data or 'prompt' not in data:
+            return jsonify({"error": "No prompt provided"}), 400
+        prompt = data['prompt']
+        print(f"Received prompt: {prompt}")
+        # Set parameters
+        batch_size = 1
+        guidance_scale = 15.0
+        # Generate latents with the text-to-3D model
+        latents = sample_latents(
+            batch_size=batch_size,
+            model=model,
+            diffusion=diffusion,
+            guidance_scale=guidance_scale,
+            model_kwargs=dict(texts=[prompt] * batch_size),
+            progress=True,
+            clip_denoised=True,
+            use_fp16=True,
+            use_karras=True,
+            karras_steps=64,
+            sigma_min=1e-3,
+            sigma_max=160,
+            s_churn=0,
+        )
+        # Generate a unique filename
+        filename = f"outputs/{uuid.uuid4()}"
+        # Convert latent to mesh
+        t0 = time.time()
+        mesh = decode_latent_mesh(xm, latents[0]).tri_mesh()
+        print(f"Mesh decoded in {time.time() - t0:.2f} seconds")
+        # Save as GLB
+        glb_path = f"{filename}.glb"
+        mesh.write_glb(glb_path)
+        # Save as OBJ
+        obj_path = f"{filename}.obj"
+        with open(obj_path, 'w') as f:
+            mesh.write_obj(f)
+        # Return paths to the generated files
+        return jsonify({
+            "success": True,
+            "message": "3D model generated successfully",
+            "glb_url": f"/download/{os.path.basename(glb_path)}",
+            "obj_url": f"/download/{os.path.basename(obj_path)}"
+        })
     except Exception as e:
+        print(f"Error: {str(e)}")
         return jsonify({"error": str(e)}), 500
+@app.route('/download/<filename>', methods=['GET'])
+def download_file(filename):
+    try:
+        return send_file(f"outputs/{filename}", as_attachment=True)
+    except Exception as e:
+        return jsonify({"error": str(e)}), 404
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860, debug=True)