Update mesh.py

mesh.py

@@ -42,6 +42,7 @@ def get_mesh(model, image: Image.Image, keep_edges=True):
         triangles = create_triangles(image.shape[0], image.shape[1])
     else:
         triangles = create_triangles(image.shape[0], image.shape[1], mask=~depth_edges_mask(depth))
+
     colors = image.reshape(-1, 3)
     mesh = trimesh.Trimesh(vertices=verts, faces=triangles, vertex_colors=colors)
 
@@ -57,4 +58,4 @@ def mesh_interface(model):
             inputs=[gr.Image(label="Input Image", type='pil'), gr.Checkbox(label="Keep occlusion edges", value=True)]
         outputs=gr.Model3D(label="3D Mesh", clear_color=[1.0, 1.0, 1.0, 1.0])
     generate_btn = gr.Button(value="Generate")
-    generate_btn.click(partial(get_mesh, model), inputs=inputs, outputs=outputs, api_name="generate_mesh")
+    generate_btn.click(partial(get_mesh, model.to(device)), inputs=inputs, outputs=outputs, api_name="generate_mesh")