Update app.py

app.py

@@ -15,7 +15,8 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from diffusers import StableFast3DPipeline
+from tsr.system import TripoSR
+from tsr.utils import remove_background, resize_foreground
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 torch.set_default_device("cpu")
@@ -42,12 +43,12 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 
 processing_jobs = {}
-sf3d_pipeline = None
+triposr_model = None
 model_loaded = False
 model_loading = False
 
 TIMEOUT_SECONDS = 300
-MAX_DIMENSION = 512  # Stable-Fast-3D uses 512x512 inputs
+MAX_DIMENSION = 512  # TripoSR uses 512x512 inputs
 
 class TimeoutError(Exception):
     pass
@@ -89,26 +90,30 @@ def preprocess_image(image_path):
             if img.mode == 'RGBA':
                 img = img.convert('RGB')
             img = img.resize((512, 512), Image.LANCZOS)
-            return img
+            img_array = np.array(img) / 255.0
+            img_array = remove_background(img_array)
+            img_array = resize_foreground(img_array, 0.85)
+            img_array = np.clip(img_array, 0, 1) * 255
+            return Image.fromarray(img_array.astype(np.uint8))
     except Exception as e:
         raise Exception(f"Error preprocessing image: {str(e)}")
 
 def load_model():
-    global sf3d_pipeline, model_loaded, model_loading
+    global triposr_model, model_loaded, model_loading
     
     if model_loaded:
-        return sf3d_pipeline
+        return triposr_model
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return sf3d_pipeline
+        return triposr_model
     
     try:
         model_loading = True
-        print("Loading Stable-Fast-3D...")
+        print("Loading TripoSR...")
         
-        model_name = "stabilityai/stable-fast-3d"
+        model_name = "tripo3d/triposr"
         
         max_retries = 3
         retry_delay = 5
@@ -128,16 +133,15 @@ def load_model():
                 else:
                     raise
         
-        sf3d_pipeline = StableFast3DPipeline.from_pretrained(
+        triposr_model = TripoSR.from_pretrained(
             model_name,
             cache_dir=CACHE_DIR,
-            torch_dtype=torch.float32,
+            device="cpu",
         )
-        sf3d_pipeline.to("cpu")
         
         model_loaded = True
-        print("Stable-Fast-3D loaded successfully on CPU")
-        return sf3d_pipeline
+        print("TripoSR loaded successfully on CPU")
+        return triposr_model
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -148,21 +152,17 @@ def load_model():
 
 def generate_3d_model(image, detail_level):
     try:
-        num_steps = {'low': 20, 'medium': 30, 'high': 50}
-        steps = num_steps[detail_level]
+        chunk_size = {'low': 4096, 'medium': 8192, 'high': 16384}
+        chunk = chunk_size[detail_level]
         
         with torch.no_grad():
-            result = sf3d_pipeline(
-                image,
-                num_inference_steps=steps,
-                normal_num_inference_steps=steps // 2,
-                guidance_scale=7.0,
-            )
+            scene_codes = triposr_model(image, device="cpu")
+            meshes = triposr_model.mesher(scene_codes, chunk_size=chunk)
         
-        mesh = result.trimesh_meshes[0]
+        mesh = meshes[0]
         vertices = np.array(mesh.vertices)
         faces = np.array(mesh.faces)
-        vertex_colors = np.array(mesh.visual.vertex_colors) if mesh.visual.vertex_colors is not None else None
+        vertex_colors = np.array(mesh.vertex_colors) if hasattr(mesh, 'vertex_colors') and mesh.vertex_colors is not None else None
         
         trimesh_mesh = trimesh.Trimesh(
             vertices=vertices,
@@ -180,7 +180,7 @@ def generate_3d_model(image, detail_level):
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "Stable-Fast-3D",
+        "model": "TripoSR",
         "device": "cpu"
     }), 200
 
@@ -268,7 +268,7 @@ def convert_image_to_3d():
             processing_jobs[job_id]['progress'] = 10
             
             try:
-                pipeline = load_model()
+                model = load_model()
                 processing_jobs[job_id]['progress'] = 30
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -294,7 +294,7 @@ def convert_image_to_3d():
                 file_path = os.path.join(output_dir, f"model.{output_format}")
                 mesh.export(file_path, file_type=output_format)
                 
-                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                processing_jobs[job_id]['result_url'] = f"/download/{job_id.ConcurrentHashMap}"
                 processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
                 
                 processing_jobs[job_id]['status'] = 'completed'
@@ -359,7 +359,7 @@ def preview_model(job_id):
         else:
             return send_file(file_path, mimetype='text/plain')
     
-    return jsonify({"error": "Model not found"}), 404
+    return jsonify({"error": "File not found"}), 404
 
 def cleanup_old_jobs():
     current_time = time.time()
@@ -420,7 +420,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (Stable-Fast-3D)",
+        "message": "Image to 3D API (TripoSR)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -432,7 +432,7 @@ def index():
             "output_format": "glb or obj",
             "detail_level": "low, medium, or high"
         },
-        "description": "Creates 3D models from 2D images using Stable-Fast-3D."
+        "description": "Creates 3D models from 2D images using TripoSR."
     }), 200
 
 if __name__ == '__main__':