Update app.py

app.py

@@ -4,6 +4,7 @@ import trimesh
 import traceback
 import numpy as np
 import gradio as gr
+from functools import partial
 from multiprocessing import Process, Queue
 
 import torch
@@ -148,113 +149,124 @@ def py_string_to_mesh_file_safe(py_string, mesh_path):
     if process.is_alive():
         process.terminate()
         process.join()
-        raise RuntimeError('Process is alive after 3 seconds')
+        raise gr.Error('Process is alive after 3 seconds')
     
     if not queue.empty():
-        raise RuntimeError(queue.get())
-
-
-def run_gpu(model, input_ids, attention_mask, point_cloud, pad_token_id):
-    if torch.cuda.is_available():
-        model = model.cuda()
-    with torch.no_grad():
-        batch_ids = model.generate(
-            input_ids=torch.tensor(input_ids).unsqueeze(0).to(model.device),
-            attention_mask=torch.tensor(attention_mask).unsqueeze(0).to(model.device),
-            point_cloud=torch.tensor(point_cloud.astype(np.float32)).unsqueeze(0).to(model.device),
-            max_new_tokens=768,
-            pad_token_id=pad_token_id).cpu()
-    return batch_ids
-
-
-def run_test(in_mesh_path, seed, results):
-    mesh = trimesh.load(in_mesh_path)
-    mesh.apply_translation(-(mesh.bounds[0] + mesh.bounds[1]) / 2.0)
-    mesh.apply_scale(2.0 / max(mesh.extents))
-    np.random.seed(seed)
-    point_cloud = mesh_to_point_cloud(mesh)
-    
-    pcd_path = '/tmp/pcd.obj'
-    trimesh.points.PointCloud(point_cloud[:, :3]).export(pcd_path)
-    results.append(pcd_path)
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        'Qwen/Qwen2-1.5B',
-        pad_token='<|im_end|>',
-        padding_side='left')
-    model = CADRecode.from_pretrained(
-        'filapro/cad-recode',
-        torch_dtype='auto').eval()
-
-    input_ids = [tokenizer.pad_token_id] * len(point_cloud) + [tokenizer('<|im_start|>')['input_ids'][0]]
-    attention_mask = [-1] * len(point_cloud) + [1]
-    batch_ids = run_gpu(model, input_ids, attention_mask, point_cloud, tokenizer.pad_token_id)
-    py_string = tokenizer.batch_decode(batch_ids)[0]
-    begin = py_string.find('<|im_start|>') + 12
-    end = py_string.find('<|endoftext|>')
-    py_string = py_string[begin: end]
-    results.append(py_string)
-
-    out_mesh_path = '/tmp/mesh.stl'
-    py_string_to_mesh_file_safe(py_string, out_mesh_path)
-    results.append(out_mesh_path)
+        raise gr.Error(queue.get())
+
+
+def run_point_cloud(in_mesh_path, seed):
+    try:
+        mesh = trimesh.load(in_mesh_path)
+        mesh.apply_translation(-(mesh.bounds[0] + mesh.bounds[1]) / 2.0)
+        mesh.apply_scale(2.0 / max(mesh.extents))
+        np.random.seed(seed)
+        point_cloud = mesh_to_point_cloud(mesh)
+        pcd_path = '/tmp/pcd.obj'
+        trimesh.points.PointCloud(point_cloud[:, :3]).export(pcd_path)
+        return point_cloud, pcd_path
+    except:
+        raise gr.Error(traceback.format_exc())
 
 
 @spaces.GPU(duration=20)
-def run_test_safe(in_mesh_path, seed):
-    results, log = list(), str()
+def run_cad_recode(point_cloud):
     try:
-        run_test(in_mesh_path, seed, results)
+        input_ids = [tokenizer.pad_token_id] * len(point_cloud) + [tokenizer('<|im_start|>')['input_ids'][0]]
+        attention_mask = [-1] * len(point_cloud) + [1]
+        if torch.cuda.is_available():
+            model = cad_recode.cuda()
+        with torch.no_grad():
+            batch_ids = cad_recode.generate(
+                input_ids=torch.tensor(input_ids).unsqueeze(0).to(model.device),
+                attention_mask=torch.tensor(attention_mask).unsqueeze(0).to(model.device),
+                point_cloud=torch.tensor(point_cloud.astype(np.float32)).unsqueeze(0).to(model.device),
+                max_new_tokens=768,
+                pad_token_id=tokenizer.pad_token_id).cpu()
+        py_string = tokenizer.batch_decode(batch_ids)[0]
+        begin = py_string.find('<|im_start|>') + 12
+        end = py_string.find('<|endoftext|>')
+        py_string = py_string[begin: end]
+        return py_string, py_string
     except:
-        log += 'Status: FAILED\n' + traceback.format_exc()
-    return results + [None] * (3 - len(results)) + [log]
-    
+        raise gr.Error(traceback.format_exc())
 
-os.environ['TOKENIZERS_PARALLELISM'] = 'False'
 
+def run_mesh(py_string):
+    try:
+        out_mesh_path = '/tmp/mesh.stl'
+        py_string_to_mesh_file_safe(py_string, out_mesh_path)
+        return out_mesh_path
+    except:
+        raise gr.Error(traceback.format_exc())
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        gr.Markdown('## CAD-Recode Demo\n'
-                    'Upload mesh or select from examples and press Run! Mesh ⇾ 256 points ⇾ Python code by CAD-Recode ⇾ CAD model.')
 
-    with gr.Row(equal_height=True):
-        in_model = gr.Model3D(label='1. Input Mesh', interactive=True)
-        point_model = gr.Model3D(label='2. Sampled Point Cloud', display_mode='point_cloud', interactive=False)
-        out_model = gr.Model3D(
-            label='4. Result CAD Model', interactive=False
+def run():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## CAD-Recode Demo\n'
+                        'Upload mesh or select from examples and press Run! Mesh ⇾ 256 points ⇾ Python code by CAD-Recode ⇾ CAD model.')
+
+        with gr.Row(equal_height=True):
+            in_model = gr.Model3D(label='1. Input Mesh', interactive=True)
+            point_model = gr.Model3D(label='2. Sampled Point Cloud', display_mode='point_cloud', interactive=False)
+            out_model = gr.Model3D(
+                label='4. Result CAD Model', interactive=False
+            )
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    seed_slider = gr.Slider(label='Random Seed', value=42, interactive=True)
+                with gr.Row():
+                    gr.Examples(
+                        examples=[
+                            ['./data/49215_5368e45e_0000.stl', 42],
+                            ['./data/00882236.stl', 6],
+                            ['./data/User Library-engrenage.stl', 18],
+                            ['./data/00010900.stl', 42],
+                            ['./data/21492_8bd34fc1_0008.stl', 42],
+                            ['./data/00375556.stl', 53],  # todo: 96?
+                            ['./data/49121_adb01620_0000.stl', 42]],
+                        example_labels=[
+                            'fusion360_table1', 'deepcad_star', 'cc3d_gear', 'deepcad_barrels',
+                            'fusion360_gear', 'deepcad_house', 'fusion360_table2'],
+                        inputs=[in_model, seed_slider],
+                        cache_examples=False)
+                with gr.Row():
+                    run_button = gr.Button('Run')
+
+            with gr.Column():
+                out_code = gr.Code(language='python', label='3. Generated Python Code', wrap_lines=True, interactive=False)
+            
+            with gr.Column():
+                pass
+
+        state = gr.State()
+        run_button.click(
+            run_point_cloud,
+            inputs=[in_model, seed_slider],
+            outputs=[state, point_model]
+        ).success(
+            run_cad_recode,
+            inputs=[state],
+            outputs=[state, out_code]
+        ).success(
+            run_mesh,
+            inputs=[state],
+            outputs=[out_model]
         )
-    
-    with gr.Row():
-        with gr.Column():
-            with gr.Row():
-                seed_slider = gr.Slider(label='Random Seed', value=42, interactive=True)
-            with gr.Row():
-                _ = gr.Examples(
-                    examples=[
-                        ['./data/49215_5368e45e_0000.stl', 42],
-                        ['./data/00882236.stl', 6],
-                        ['./data/User Library-engrenage.stl', 18],
-                        ['./data/00010900.stl', 42],
-                        ['./data/21492_8bd34fc1_0008.stl', 42],
-                        ['./data/00375556.stl', 96],
-                        ['./data/49121_adb01620_0000.stl', 42]],
-                    example_labels=[
-                        'fusion360_table1', 'deepcad_star', 'cc3d_gear', 'deepcad_barrels',
-                        'fusion360_gear', 'deepcad_house', 'fusion360_table2'],
-                    inputs=[in_model, seed_slider],
-                    cache_examples=False)
-            with gr.Row():
-                run_button = gr.Button('Run')
-                _ = gr.LoginButton()
-
-        with gr.Column():
-            out_code = gr.Code(language='python', label='3. Generated Python Code', wrap_lines=True, interactive=False)
-
-        with gr.Column():
-            log_textbox = gr.Textbox(label='Log', placeholder='Status: OK', interactive=False)
-    
-    run_button.click(
-        run_test_safe, inputs=[in_model, seed_slider], outputs=[point_model, out_code, out_model, log_textbox])
 
-demo.launch()
+    demo.launch(show_error=True)
+
+
+tokenizer = AutoTokenizer.from_pretrained(
+    'Qwen/Qwen2-1.5B',
+    pad_token='<|im_end|>',
+    padding_side='left')
+cad_recode = CADRecode.from_pretrained(
+    'filapro/cad-recode',
+    torch_dtype='auto').eval()
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'False'
+run()