Update app.py

app.py

@@ -8,23 +8,20 @@ from flask import Flask, request, jsonify, send_file, Response, stream_with_cont
 from werkzeug.utils import secure_filename
 from PIL import Image
 import io
-import zipfile
 import uuid
 import traceback
 from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-import cv2
-import pymeshlab
-from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline
+from instantmesh import InstantMesh
 
 # Force CPU usage
-os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Disable GPU detection
-torch.set_default_device("cpu")  # Set CPU as default device
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+torch.set_default_device("cpu")
 
 app = Flask(__name__)
-CORS(app)  # Enable CORS for all routes
+CORS(app)
 
 # Configure directories
 UPLOAD_FOLDER = '/tmp/uploads'
@@ -43,19 +40,19 @@ os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
 os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 
 # Job tracking dictionary
 processing_jobs = {}
 
 # Global model variables
-hunyuan_pipeline = None
+instant_mesh = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 600  # 10 minutes max for Hunyuan3D-2mini on CPU
-MAX_DIMENSION = 256    # Reduced for CPU memory constraints
+TIMEOUT_SECONDS = 300  # 5 minutes max for InstantMesh on CPU
+MAX_DIMENSION = 256
 
 # TimeoutError for handling timeouts
 class TimeoutError(Exception):
@@ -94,12 +91,10 @@ def process_with_timeout(function, args, timeout):
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-# Simplified image preprocessing for Hunyuan3D-2mini
+# Simplified image preprocessing
 def preprocess_image(image_path):
     with Image.open(image_path) as img:
         img = img.convert("RGB")
-        
-        # Resize to smaller dimensions for CPU
         if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
             if img.width > img.height:
                 new_width = MAX_DIMENSION
@@ -108,25 +103,24 @@ def preprocess_image(image_path):
                 new_height = MAX_DIMENSION
                 new_width = int(img.width * (MAX_DIMENSION / img.height))
             img = img.resize((new_width, new_height), Image.LANCZOS)
-        
         return img
 
 def load_model():
-    global hunyuan_pipeline, model_loaded, model_loading
+    global instant_mesh, model_loaded, model_loading
     
     if model_loaded:
-        return hunyuan_pipeline
+        return instant_mesh
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return hunyuan_pipeline
+        return instant_mesh
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        model_name = "tencent/Hunyuan3D-2mini"
+        model_name = "BAAI/InstantMesh"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -147,19 +141,18 @@ def load_model():
                 else:
                     raise
         
-        # Load Hunyuan3D-2mini pipeline
-        hunyuan_pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
+        # Load InstantMesh
+        instant_mesh = InstantMesh.from_pretrained(
             model_name,
-            subfolder="hunyuan3d-dit-v2-mini",
             use_safetensors=True,
             torch_dtype=torch.float16,
             cache_dir=CACHE_DIR,
-            device_map="cpu"  # Explicitly set to CPU
+            device="cpu"
         )
         
         model_loaded = True
         print("Model loaded successfully on CPU")
-        return hunyuan_pipeline
+        return instant_mesh
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -168,20 +161,11 @@ def load_model():
     finally:
         model_loading = False
 
-# Optimize mesh for Unity
-def optimize_mesh(mesh_path, target_faces=10000):
-    ms = pymeshlab.MeshSet()
-    ms.load_new_mesh(mesh_path)
-    ms.meshing_decimation_quadric_edge_collapse(targetfacenum=target_faces)
-    optimized_path = mesh_path.replace(".glb", "_optimized.glb")
-    ms.save_current_mesh(optimized_path)
-    return optimized_path
-
 @app.route('/health', methods=['GET'])
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "Hunyuan3D-2mini 3D Generator",
+        "model": "InstantMesh 3D Generator",
         "device": "cpu"
     }), 200
 
@@ -239,7 +223,7 @@ def convert_image_to_3d():
         return jsonify({"error": "Invalid parameter values"}), 400
     
     if output_format not in ['glb']:
-        return jsonify({"error": "Only GLB format is supported with Hunyuan3D-2mini"}), 400
+        return jsonify({"error": "Only GLB format is supported"}), 400
     
     job_id = str(uuid.uuid4())
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
@@ -279,17 +263,19 @@ def convert_image_to_3d():
             try:
                 def generate_3d():
                     # Adjust settings based on detail level
-                    steps = {'low': 20, 'medium': 30, 'high': 40}
-                    resolution = {'low': 200, 'medium': 256, 'high': 300}
+                    steps = {'low': 50, 'medium': 75, 'high': 100}
+                    resolution = {'low': 128, 'medium': 256, 'high': 512}
                     
+                    # Convert PIL image to numpy for InstantMesh
+                    img_array = np.array(image)
+                    
+                    # Generate mesh
                     mesh = model(
-                        image=image,
-                        num_inference_steps=steps[detail_level],
-                        octree_resolution=resolution[detail_level],
-                        num_chunks=10000,
-                        generator=torch.manual_seed(12345),
-                        output_type="trimesh"
-                    )[0]
+                        image=img_array,
+                        num_steps=steps[detail_level],
+                        resolution=resolution[detail_level],
+                        seed=12345
+                    )
                     return mesh
                 
                 mesh, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
@@ -304,13 +290,10 @@ def convert_image_to_3d():
                         
                 processing_jobs[job_id]['progress'] = 80
                 
-                # Export and optimize
+                # Export as GLB
                 glb_path = os.path.join(output_dir, "model.glb")
                 mesh.export(glb_path, file_type='glb')
                 
-                # Optimize for Unity
-                optimized_path = optimize_mesh(glb_path)
-                
                 processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
                 processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
                 
@@ -353,7 +336,7 @@ def download_model(job_id):
         return jsonify({"error": "Model not found or processing not complete"}), 404
     
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
-    glb_path = os.path.join(output_dir, "model_optimized.glb")
+    glb_path = os.path.join(output_dir, "model.glb")
     
     if os.path.exists(glb_path):
         return send_file(glb_path, as_attachment=True, download_name="model.glb")
@@ -366,7 +349,7 @@ def preview_model(job_id):
         return jsonify({"error": "Model not found or processing not complete"}), 404
     
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
-    glb_path = os.path.join(output_dir, "model_optimized.glb")
+    glb_path = os.path.join(output_dir, "model.glb")
     
     if os.path.exists(glb_path):
         return send_file(glb_path, mimetype='model/gltf-binary')
@@ -414,7 +397,7 @@ def model_info(job_id):
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
     model_stats = {}
     
-    glb_path = os.path.join(output_dir, "model_optimized.glb")
+    glb_path = os.path.join(output_dir, "model.glb")
     if os.path.exists(glb_path):
         model_stats['model_size'] = os.path.getsize(glb_path)
     
@@ -431,7 +414,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (Hunyuan3D-2mini)",
+        "message": "Image to 3D API (InstantMesh)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -443,7 +426,7 @@ def index():
             "output_format": "glb",
             "detail_level": "low, medium, or high - controls mesh detail"
         },
-        "description": "This API creates full 3D models from 2D images using Hunyuan3D-2mini"
+        "description": "This API creates full 3D models from 2D images using InstantMesh"
     }), 200
 
 if __name__ == '__main__':