Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import os
 import torch
 import time
@@ -16,11 +17,11 @@ from flask_cors import CORS
 import numpy as np
 import trimesh
 import cv2
-from tsr.system import TSR  # Updated import
-import torchvision.transforms as T
+import pymeshlab
+from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline
 
 app = Flask(__name__)
-CORS(app)
+CORS(app)  # Enable CORS for all routes
 
 # Configure directories
 UPLOAD_FOLDER = '/tmp/uploads'
@@ -28,12 +29,12 @@ RESULTS_FOLDER = '/tmp/results'
 CACHE_DIR = '/tmp/huggingface'
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
-# Create directories
+# Create necessary directories
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(RESULTS_FOLDER, exist_ok=True)
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Set Hugging Face cache
+# Set Hugging Face cache environment variables
 os.environ['HF_HOME'] = CACHE_DIR
 os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
 os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
@@ -41,22 +42,23 @@ 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
 
-# Job tracking
+# Job tracking dictionary
 processing_jobs = {}
 
 # Global model variables
-u2net_model = None
-triposr_model = None
+hunyuan_pipeline = None
 model_loaded = False
 model_loading = False
 
-# Configuration
-TIMEOUT_SECONDS = 240  # 4 minutes max
-MAX_DIMENSION = 512    # Max image dimension
+# Configuration for processing
+TIMEOUT_SECONDS = 600  # 10 minutes max for Hunyuan3D-2mini on CPU
+MAX_DIMENSION = 256    # Reduced for CPU memory constraints
 
+# TimeoutError for handling timeouts
 class TimeoutError(Exception):
     pass
 
+# Thread-safe timeout implementation
 def process_with_timeout(function, args, timeout):
     result = [None]
     error = [None]
@@ -89,11 +91,12 @@ 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
 def preprocess_image(image_path):
     with Image.open(image_path) as img:
         img = img.convert("RGB")
         
-        # Resize if too large
+        # 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
@@ -103,62 +106,28 @@ def preprocess_image(image_path):
                 new_width = int(img.width * (MAX_DIMENSION / img.height))
             img = img.resize((new_width, new_height), Image.LANCZOS)
         
-        # Apply adaptive histogram equalization
-        img_array = np.array(img)
-        if len(img_array.shape) == 3 and img_array.shape[2] == 3:
-            lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
-            l, a, b = cv2.split(lab)
-            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-            cl = clahe.apply(l)
-            enhanced_lab = cv2.merge((cl, a, b))
-            img_array = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
-            img = Image.fromarray(img_array)
-            
         return img
 
-def remove_background(image):
-    global u2net_model
-    if u2net_model is None:
-        # Dynamically import U2NET to avoid circular import issues
-        from u2net import U2NET
-        u2net_model = U2NET()
-        u2net_model.load_state_dict(torch.load('u2net.pth', map_location='cpu'))
-        u2net_model.eval()
-        u2net_model.to('cpu')
-    
-    img_array = np.array(image)
-    img_tensor = T.ToTensor()(image.resize((320, 320))).unsqueeze(0)
-    
-    with torch.no_grad():
-        d1, *_ = u2net_model(img_tensor)
-        pred = d1[:, 0, :, :]
-        pred = (pred - pred.min()) / (pred.max() - pred.min())
-        mask = (pred > 0.5).float().squeeze().numpy()
-    
-    mask_img = Image.fromarray((mask * 255).astype('uint8')).resize(image.size)
-    mask_array = np.array(mask_img)[:, :, np.newaxis] / 255
-    result = img_array * mask_array + (1 - mask_array) * 255  # White background
-    return Image.fromarray(result.astype('uint8'))
-
 def load_model():
-    global triposr_model, model_loaded, model_loading
+    global hunyuan_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return triposr_model
+        return hunyuan_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return triposr_model
+        return hunyuan_pipeline
     
     try:
         model_loading = True
-        print("Loading TripoSR model...")
+        print("Starting model loading...")
         
-        model_name = "stabilityai/TripoSR"
+        model_name = "tencent/Hunyuan3D-2mini"
+        
+        # Download model with retry mechanism
         max_retries = 3
         retry_delay = 5
-        
         for attempt in range(max_retries):
             try:
                 snapshot_download(
@@ -169,23 +138,27 @@ def load_model():
                 break
             except Exception as e:
                 if attempt < max_retries - 1:
-                    print(f"Download attempt {attempt+1} failed: {str(e)}. Retrying...")
+                    print(f"Download attempt {attempt+1} failed: {str(e)}. Retrying in {retry_delay} seconds...")
                     time.sleep(retry_delay)
                     retry_delay *= 2
                 else:
                     raise
         
-        # Initialize TSR model
-        triposr_model = TSR.from_pretrained(
+        # Load Hunyuan3D-2mini pipeline
+        hunyuan_pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
             model_name,
-            torch_dtype=torch.float32,
-            device="cpu",
+            subfolder="hunyuan3d-dit-v2-mini",
+            use_safetensors=True,
+            torch_dtype=torch.float16,
             cache_dir=CACHE_DIR
         )
         
+        # Move to CPU
+        hunyuan_pipeline.to("cpu")
+        
         model_loaded = True
-        print("TripoSR model loaded successfully on CPU")
-        return triposr_model
+        print("Model loaded successfully on CPU")
+        return hunyuan_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -194,27 +167,20 @@ def load_model():
     finally:
         model_loading = False
 
-def optimize_mesh(mesh, detail_level='medium'):
-    # Simplify mesh based on detail level
-    if detail_level == 'high':
-        target_faces = 50000
-    elif detail_level == 'medium':
-        target_faces = 30000
-    else:
-        target_faces = 15000
-    
-    if len(mesh.faces) > target_faces:
-        mesh = mesh.simplify_quadric_decimation(target_faces)
-    
-    # Fix normals
-    mesh.fix_normals()
-    return mesh
+# 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": "TripoSR 3D Model Generator",
+        "model": "Hunyuan3D-2mini 3D Generator",
         "device": "cpu"
     }), 200
 
@@ -226,6 +192,7 @@ def progress(job_id):
             return
             
         job = processing_jobs[job_id]
+        
         yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
         
         last_progress = job['progress']
@@ -234,8 +201,10 @@ def progress(job_id):
             if job['progress'] != last_progress:
                 yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
                 last_progress = job['progress']
+            
             time.sleep(0.5)
             check_count += 1
+            
             if check_count > 60:
                 if 'thread_alive' in job and not job['thread_alive']():
                     job['status'] = 'error'
@@ -260,17 +229,16 @@ def convert_image_to_3d():
         return jsonify({"error": "No image selected"}), 400
     
     if not allowed_file(file.filename):
-        return jsonify({"error": f"File type not allowed: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
+        return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
     
     try:
         output_format = request.form.get('output_format', 'glb').lower()
         detail_level = request.form.get('detail_level', 'medium').lower()
-        texture_quality = request.form.get('texture_quality', 'medium').lower()
     except ValueError:
         return jsonify({"error": "Invalid parameter values"}), 400
     
-    if output_format not in ['obj', 'glb']:
-        return jsonify({"error": "Unsupported output format: 'obj' or 'glb'"}), 400
+    if output_format not in ['glb']:
+        return jsonify({"error": "Only GLB format is supported with Hunyuan3D-2mini"}), 400
     
     job_id = str(uuid.uuid4())
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
@@ -295,30 +263,32 @@ def convert_image_to_3d():
         processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
         
         try:
-            # Preprocess image
             processing_jobs[job_id]['progress'] = 5
             image = preprocess_image(filepath)
             processing_jobs[job_id]['progress'] = 10
             
-            # Remove background
-            processing_jobs[job_id]['progress'] = 20
-            clean_image = remove_background(image)
-            processing_jobs[job_id]['progress'] = 30
-            
-            # Load TripoSR model
             try:
                 model = load_model()
-                processing_jobs[job_id]['progress'] = 40
+                processing_jobs[job_id]['progress'] = 30
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
                 processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
                 return
             
-            # Generate 3D model
             try:
                 def generate_3d():
-                    # TSR expects a PIL image
-                    mesh = model(clean_image)
+                    # Adjust settings based on detail level
+                    steps = {'low': 20, 'medium': 30, 'high': 40}
+                    resolution = {'low': 200, 'medium': 256, 'high': 300}
+                    
+                    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]
                     return mesh
                 
                 mesh, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
@@ -330,49 +300,18 @@ def convert_image_to_3d():
                         return
                     else:
                         raise error
-                
-                processing_jobs[job_id]['progress'] = 70
-                
-                # Optimize mesh
-                mesh = optimize_mesh(mesh, detail_level)
+                        
                 processing_jobs[job_id]['progress'] = 80
                 
-            except Exception as e:
-                error_details = traceback.format_exc()
-                processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
-                print(f"Error processing job {job_id}: {str(e)}")
-                print(error_details)
-                return
-            
-            # Export model
-            try:
-                if output_format == 'obj':
-                    obj_path = os.path.join(output_dir, "model.obj")
-                    mesh.export(
-                        obj_path,
-                        file_type='obj',
-                        include_normals=True,
-                        include_texture=True
-                    )
-                    zip_path = os.path.join(output_dir, "model.zip")
-                    with zipfile.ZipFile(zip_path, 'w') as zipf:
-                        zipf.write(obj_path, arcname="model.obj")
-                        mtl_path = os.path.join(output_dir, "model.mtl")
-                        if os.path.exists(mtl_path):
-                            zipf.write(mtl_path, arcname="model.mtl")
-                        texture_path = os.path.join(output_dir, "model.png")
-                        if os.path.exists(texture_path):
-                            zipf.write(texture_path, arcname="model.png")
-                    
-                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                # Export and optimize
+                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)
                 
-                elif output_format == 'glb':
-                    glb_path = os.path.join(output_dir, "model.glb")
-                    mesh.export(glb_path, file_type='glb')
-                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
                 
                 processing_jobs[job_id]['status'] = 'completed'
                 processing_jobs[job_id]['progress'] = 100
@@ -381,9 +320,10 @@ def convert_image_to_3d():
             except Exception as e:
                 error_details = traceback.format_exc()
                 processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Error exporting model: {str(e)}"
-                print(f"Error exporting model for job {job_id}: {str(e)}")
+                processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
+                print(f"Error processing job {job_id}: {str(e)}")
                 print(error_details)
+                return
             
             if os.path.exists(filepath):
                 os.remove(filepath)
@@ -396,6 +336,7 @@ def convert_image_to_3d():
             processing_jobs[job_id]['error'] = f"{str(e)}\n{error_details}"
             print(f"Error processing job {job_id}: {str(e)}")
             print(error_details)
+            
             if os.path.exists(filepath):
                 os.remove(filepath)
     
@@ -411,16 +352,10 @@ 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)
-    output_format = processing_jobs[job_id].get('output_format', 'glb')
+    glb_path = os.path.join(output_dir, "model_optimized.glb")
     
-    if output_format == 'obj':
-        zip_path = os.path.join(output_dir, "model.zip")
-        if os.path.exists(zip_path):
-            return send_file(zip_path, as_attachment=True, download_name="model.zip")
-    else:
-        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")
+    if os.path.exists(glb_path):
+        return send_file(glb_path, as_attachment=True, download_name="model.glb")
     
     return jsonify({"error": "File not found"}), 404
 
@@ -430,16 +365,10 @@ 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)
-    output_format = processing_jobs[job_id].get('output_format', 'glb')
-
-    if output_format == 'obj':
-        obj_path = os.path.join(output_dir, "model.obj")
-        if os.path.exists(obj_path):
-            return send_file(obj_path, mimetype='model/obj')
-    else:
-        glb_path = os.path.join(output_dir, "model.glb")
-        if os.path.exists(glb_path):
-            return send_file(glb_path, mimetype='model/gltf-binary')
+    glb_path = os.path.join(output_dir, "model_optimized.glb")
+    
+    if os.path.exists(glb_path):
+        return send_file(glb_path, mimetype='model/gltf-binary')
     
     return jsonify({"error": "Model file not found"}), 404
 
@@ -461,6 +390,7 @@ def cleanup_old_jobs():
                 shutil.rmtree(output_dir)
         except Exception as e:
             print(f"Error cleaning up job {job_id}: {str(e)}")
+        
         if job_id in processing_jobs:
             del processing_jobs[job_id]
     
@@ -483,17 +413,9 @@ def model_info(job_id):
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
     model_stats = {}
     
-    if job['output_format'] == 'obj':
-        obj_path = os.path.join(output_dir, "model.obj")
-        zip_path = os.path.join(output_dir, "model.zip")
-        if os.path.exists(obj_path):
-            model_stats['obj_size'] = os.path.getsize(obj_path)
-        if os.path.exists(zip_path):
-            model_stats['package_size'] = os.path.getsize(zip_path)
-    else:
-        glb_path = os.path.join(output_dir, "model.glb")
-        if os.path.exists(glb_path):
-            model_stats['model_size'] = os.path.getsize(glb_path)
+    glb_path = os.path.join(output_dir, "model_optimized.glb")
+    if os.path.exists(glb_path):
+        model_stats['model_size'] = os.path.getsize(glb_path)
     
     return jsonify({
         "status": job['status'],
@@ -508,7 +430,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "TripoSR Image to 3D API",
+        "message": "Image to 3D API (Hunyuan3D-2mini)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -517,14 +439,14 @@ def index():
             "/model-info/<job_id>"
         ],
         "parameters": {
-            "output_format": "obj or glb",
-            "detail_level": "low, medium, or high - controls mesh density",
-            "texture_quality": "low, medium, or high - controls texture quality"
+            "output_format": "glb",
+            "detail_level": "low, medium, or high - controls mesh detail"
         },
-        "description": "Creates full 3D models from 2D images with background removal"
+        "description": "This API creates full 3D models from 2D images using Hunyuan3D-2mini"
     }), 200
 
 if __name__ == '__main__':
     cleanup_old_jobs()
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)
+```