Update app.py

app.py

@@ -1,10 +1,8 @@
 import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
-
 import os
 import shutil
-import trimesh  # New import
 os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
 import torch
@@ -15,164 +13,161 @@ from PIL import Image
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
-from scipy.spatial import ConvexHull  # New import
-
-# [Previous imports and constants remain the same...]
-
-def optimize_building_mesh(mesh, angle_threshold=15, planar_threshold=0.02):
-    """
-    Optimize a building mesh by preserving architectural features while reducing complexity.
-    """
-    # Convert vertices to numpy array for processing
-    vertices = np.array(mesh.vertices)
-    faces = np.array(mesh.faces)
-    
-    # 1. Detect planar surfaces
-    normals = mesh.face_normals
-    planar_groups = []
-    processed = set()
-    
+from scipy.spatial.transform import Rotation
+
+# Constants
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+os.makedirs(TMP_DIR, exist_ok=True)
+
+# Session Management
+def start_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+
+def end_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    shutil.rmtree(user_dir)
+
+# Utility Functions
+def preprocess_image(image: Image.Image) -> Image.Image:
+    return pipeline.preprocess_image(image)
+
+def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
+    images = [image[0] for image in images]
+    return [pipeline.preprocess_image(image) for image in images]
+
+def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
+    return {
+        'gaussian': {
+            **gs.init_params,
+            '_xyz': gs._xyz.cpu().numpy(),
+            '_features_dc': gs._features_dc.cpu().numpy(),
+            '_scaling': gs._scaling.cpu().numpy(),
+            '_rotation': gs._rotation.cpu().numpy(),
+            '_opacity': gs._opacity.cpu().numpy(),
+        },
+        'mesh': {
+            'vertices': mesh.vertices.cpu().numpy(),
+            'faces': mesh.faces.cpu().numpy(),
+        },
+    }
+
+def unpack_state(state: dict) -> Tuple[Gaussian, MeshExtractResult]:
+    gs = Gaussian(
+        aabb=state['gaussian']['aabb'],
+        sh_degree=state['gaussian']['sh_degree'],
+        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+        scaling_bias=state['gaussian']['scaling_bias'],
+        opacity_bias=state['gaussian']['opacity_bias'],
+        scaling_activation=state['gaussian']['scaling_activation'],
+    )
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+
+    mesh = MeshExtractResult(
+        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    )
+
+    return gs, mesh
+
+def convert_to_poly_surfaces(mesh: MeshExtractResult, normal_threshold: float = 0.95) -> MeshExtractResult:
+    vertices = mesh.vertices.cpu().numpy()
+    faces = mesh.faces.cpu().numpy()
+
+    v0 = vertices[faces[:, 0]]
+    v1 = vertices[faces[:, 1]]
+    v2 = vertices[faces[:, 2]]
+    normals = np.cross(v1 - v0, v2 - v0)
+    norms = np.linalg.norm(normals, axis=1)[:, None]
+    norms[norms == 0] = 1e-10
+    normals = normals / norms
+
+    groups = []
+    used = set()
+
     for i in range(len(faces)):
-        if i in processed:
+        if i in used:
             continue
-        
-        # Find connected faces with similar normals
-        similar_faces = {i}
-        stack = [i]
-        while stack:
-            current = stack.pop()
-            neighbors = mesh.face_adjacency[mesh.face_adjacency[:,0] == current][:,1]
-            for n in neighbors:
-                if n not in processed:
-                    angle = np.arccos(np.dot(normals[current], normals[n])) * 180 / np.pi
-                    if angle < angle_threshold:
-                        similar_faces.add(n)
-                        stack.append(n)
-                        processed.add(n)
-        
-        if len(similar_faces) > 0:
-            planar_groups.append(list(similar_faces))
-    
-    # 2. Simplify each planar group while preserving edges
-    new_vertices = []
+
+        group = {i}
+        used.add(i)
+
+        for j in range(i + 1, len(faces)):
+            if j in used:
+                continue
+
+            if np.abs(np.dot(normals[i], normals[j])) > normal_threshold:
+                if len(set(faces[i]).intersection(faces[j])) >= 2:
+                    group.add(j)
+                    used.add(j)
+
+        groups.append(list(group))
+
     new_faces = []
-    vertex_map = {}
-    
-    for group in planar_groups:
-        # Get vertices for this group
-        group_faces = faces[group]
-        group_verts = vertices[np.unique(group_faces)]
-        
-        # Find best fitting plane
-        centroid = np.mean(group_verts, axis=0)
-        _, _, vh = np.linalg.svd(group_verts - centroid)
-        normal = vh[2]
-        
-        # Project vertices to plane and simplify
-        projected = group_verts - np.dot(group_verts - centroid, normal)[:, np.newaxis] * normal
-        
-        # Create simplified convex hull for this section
-        hull = ConvexHull(projected[:,:2])
-        hull_vertices = projected[hull.vertices]
-        
-        # Add to new mesh
-        start_idx = len(new_vertices)
-        new_vertices.extend(hull_vertices)
-        
-        # Triangulate the hull
-        for i in range(1, len(hull_vertices) - 1):
-            new_faces.append([start_idx, start_idx + i, start_idx + i + 1])
-    
-    # 3. Create new optimized mesh
-    optimized_mesh = trimesh.Trimesh(
-        vertices=np.array(new_vertices),
-        faces=np.array(new_faces)
+    for group in groups:
+        if len(group) <= 2:
+            for idx in group:
+                new_faces.append(faces[idx])
+        else:
+            group_faces = faces[group]
+            vert_mask = np.zeros(len(vertices), dtype=bool)
+            vert_mask[group_faces.flatten()] = True
+
+            group_verts = vertices[vert_mask]
+            normal = normals[group[0]]
+            rot = Rotation.align_vectors([[0, 0, 1]], [normal])[0]
+            projected = rot.apply(group_verts)
+            for idx in group:
+                new_faces.append(faces[idx])
+
+    new_faces = np.array(new_faces)
+    return MeshExtractResult(
+        vertices=mesh.vertices,
+        faces=torch.tensor(new_faces, device=mesh.faces.device),
     )
-    
-    return optimized_mesh
 
-# Modify the existing extract_glb function
+# Main Functions
+@spaces.GPU
+def image_to_3d(image: Image.Image, multiimages: List[Tuple[Image.Image, str]], is_multiimage: bool, seed: int,
+                ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float,
+                slat_sampling_steps: int, multiimage_algo: Literal["multidiffusion", "stochastic"], req: gr.Request) -> Tuple[dict, str]:
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    outputs = pipeline.run(image, seed=seed, formats=["gaussian", "mesh"], preprocess_image=False,
+                           sparse_structure_sampler_params={"steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength},
+                           slat_sampler_params={"steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength})
+
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_path = os.path.join(user_dir, 'sample.mp4')
+    imageio.mimsave(video_path, video, fps=15)
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+    torch.cuda.empty_cache()
+    return state, video_path
+
 @spaces.GPU(duration=90)
-def extract_glb(
-    state: dict,
-    mesh_simplify: float,
-    texture_size: int,
-    is_building: bool,  # New parameter
-    angle_threshold: float,  # New parameter
-    planar_threshold: float,  # New parameter
-    req: gr.Request,
-) -> Tuple[str, str]:
-    """
-    Extract a GLB file from the 3D model with optional building optimization.
-    """
+def extract_glb(state: dict, mesh_simplify: float, texture_size: int, use_poly_surfaces: bool,
+                normal_threshold: float, req: gr.Request) -> Tuple[str, str]:
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh = unpack_state(state)
-    
-    if is_building:
-        # Convert to trimesh for optimization
-        trimesh_mesh = trimesh.Trimesh(
-            vertices=mesh.vertices.cpu().numpy(),
-            faces=mesh.faces.cpu().numpy()
-        )
-        
-        # Apply building-specific optimization
-        optimized_mesh = optimize_building_mesh(
-            trimesh_mesh,
-            angle_threshold=angle_threshold,
-            planar_threshold=planar_threshold
-        )
-        
-        # Convert back to original format
-        mesh.vertices = torch.tensor(optimized_mesh.vertices, device='cuda')
-        mesh.faces = torch.tensor(optimized_mesh.faces, device='cuda')
-    
+
+    if use_poly_surfaces:
+        mesh = convert_to_poly_surfaces(mesh, normal_threshold)
+
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, 'sample.glb')
     glb.export(glb_path)
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-# Modify the main UI code section
-with gr.Blocks(delete_cache=(600, 600)) as demo:
-    # [Previous UI code remains the same until GLB Extraction Settings...]
-    
-    with gr.Accordion(label="GLB Extraction Settings", open=False):
-        mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
-        texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
-        # Add new building optimization controls
-        with gr.Row():
-            is_building = gr.Checkbox(label="Enable Building Optimization", value=False)
-        with gr.Column(visible=False) as building_settings:
-            angle_threshold = gr.Slider(5, 45, label="Edge Angle Threshold", value=15, step=1)
-            planar_threshold = gr.Slider(0.01, 0.1, label="Planar Surface Threshold", value=0.02, step=0.01)
-    
-    # [Rest of the UI code remains the same until the event handlers...]
-    
-    # Add visibility toggle for building settings
-    is_building.change(
-        lambda x: gr.Column.update(visible=x),
-        inputs=[is_building],
-        outputs=[building_settings]
-    )
-    
-    # Modify the extract_glb button click handler
-    extract_glb_btn.click(
-        extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size, is_building, angle_threshold, planar_threshold],
-        outputs=[model_output, download_glb],
-    ).then(
-        lambda: gr.Button(interactive=True),
-        outputs=[download_glb],
-    )
-    
-    # [Rest of the code remains the same...]
-
+# Gradio Interface
+demo = gr.Blocks()
+# Add UI elements similar to the original code
 # Launch the Gradio app
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
-    try:
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
-    except:
-        pass
-    demo.launch()
+    demo.launch()