Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import torch
 import time
@@ -15,8 +14,8 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from tsr.system import TripoSR
-from tsr.utils import remove_background, resize_foreground
+import cv2
+from lgm.models import LGM
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 torch.set_default_device("cpu")
@@ -43,12 +42,12 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 
 processing_jobs = {}
-triposr_model = None
+lgm_model = None
 model_loaded = False
 model_loading = False
 
 TIMEOUT_SECONDS = 300
-MAX_DIMENSION = 512  # TripoSR uses 512x512 inputs
+MAX_DIMENSION = 512  # LGM uses 512x512 inputs
 
 class TimeoutError(Exception):
     pass
@@ -86,34 +85,35 @@ def allowed_file(filename):
 
 def preprocess_image(image_path):
     try:
-        with Image.open(image_path) as img:
-            if img.mode == 'RGBA':
-                img = img.convert('RGB')
-            img = img.resize((512, 512), Image.LANCZOS)
-            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))
+        img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
+        if img.shape[2] == 4:  # RGBA
+            alpha = img[:, :, 3]
+            rgb = img[:, :, :3]
+            white_bg = np.ones_like(rgb) * 255
+            mask = alpha[:, :, np.newaxis] / 255.0
+            img = rgb * mask + white_bg * (1 - mask)
+            img = img.astype(np.uint8)
+        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LANCZOS4)
+        return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
     except Exception as e:
         raise Exception(f"Error preprocessing image: {str(e)}")
 
 def load_model():
-    global triposr_model, model_loaded, model_loading
+    global lgm_model, model_loaded, model_loading
     
     if model_loaded:
-        return triposr_model
+        return lgm_model
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return triposr_model
+        return lgm_model
     
     try:
         model_loading = True
-        print("Loading TripoSR...")
+        print("Loading LGM...")
         
-        model_name = "tripo3d/triposr"
+        model_name = "large-gaussian-model/lgm"
         
         max_retries = 3
         retry_delay = 5
@@ -133,15 +133,15 @@ def load_model():
                 else:
                     raise
         
-        triposr_model = TripoSR.from_pretrained(
+        lgm_model = LGM.from_pretrained(
             model_name,
             cache_dir=CACHE_DIR,
             device="cpu",
         )
         
         model_loaded = True
-        print("TripoSR loaded successfully on CPU")
-        return triposr_model
+        print("LGM loaded successfully on CPU")
+        return lgm_model
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -152,17 +152,19 @@ def load_model():
 
 def generate_3d_model(image, detail_level):
     try:
-        chunk_size = {'low': 4096, 'medium': 8192, 'high': 16384}
-        chunk = chunk_size[detail_level]
+        resolution = {'low': 256, 'medium': 512, 'high': 1024}
+        res = resolution[detail_level]
         
         with torch.no_grad():
-            scene_codes = triposr_model(image, device="cpu")
-            meshes = triposr_model.mesher(scene_codes, chunk_size=chunk)
+            mesh_data = lgm_model.generate_mesh(
+                image,
+                resolution=res,
+                device="cpu"
+            )
         
-        mesh = meshes[0]
-        vertices = np.array(mesh.vertices)
-        faces = np.array(mesh.faces)
-        vertex_colors = np.array(mesh.vertex_colors) if hasattr(mesh, 'vertex_colors') and mesh.vertex_colors is not None else None
+        vertices = np.array(mesh_data['vertices'])
+        faces = np.array(mesh_data['faces'])
+        vertex_colors = np.array(mesh_data['vertex_colors']) if 'vertex_colors' in mesh_data else None
         
         trimesh_mesh = trimesh.Trimesh(
             vertices=vertices,
@@ -180,7 +182,7 @@ def generate_3d_model(image, detail_level):
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "TripoSR",
+        "model": "LGM",
         "device": "cpu"
     }), 200
 
@@ -294,7 +296,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.ConcurrentHashMap}"
+                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'
@@ -420,7 +422,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (TripoSR)",
+        "message": "Image to 3D API (LGM)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -432,7 +434,7 @@ def index():
             "output_format": "glb or obj",
             "detail_level": "low, medium, or high"
         },
-        "description": "Creates 3D models from 2D images using TripoSR."
+        "description": "Creates 3D models from 2D images using LGM."
     }), 200
 
 if __name__ == '__main__':