暂时完成业务逻辑

app.py

@@ -8,6 +8,7 @@ from pathlib import Path
 import uuid
 import argparse
 import torch
+import trimesh
 
 
 parser = argparse.ArgumentParser()
@@ -30,6 +31,11 @@ args.enable_flashvdm = True
 SAVE_DIR = args.cache_path
 os.makedirs(SAVE_DIR, exist_ok=True)
 
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+HTML_HEIGHT = 690
+HTML_WIDTH = 500
+
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -57,9 +63,53 @@ def gen_save_folder(max_size=200):
 
     return new_folder
 
+def export_mesh(mesh, save_folder, textured=False, type='glb'):
+    if textured:
+        path = os.path.join(save_folder, f'textured_mesh.{type}')
+    else:
+        path = os.path.join(save_folder, f'white_mesh.{type}')
+    if type not in ['glb', 'obj']:
+        mesh.export(path)
+    else:
+        mesh.export(path, include_normals=textured)
+    return path
+
+def build_model_viewer_html(save_folder, height=660, width=790, textured=False):
+    # Remove first folder from path to make relative path
+    if textured:
+        related_path = f"./textured_mesh.glb"
+        template_name = './assets/modelviewer-textured-template.html'
+        output_html_path = os.path.join(save_folder, f'textured_mesh.html')
+    else:
+        related_path = f"./white_mesh.glb"
+        template_name = './assets/modelviewer-template.html'
+        output_html_path = os.path.join(save_folder, f'white_mesh.html')
+        
+    offset = 50 if textured else 10
+    with open(os.path.join(CURRENT_DIR, template_name), 'r', encoding='utf-8') as f:
+        template_html = f.read()
+
+    with open(output_html_path, 'w', encoding='utf-8') as f:
+        template_html = template_html.replace('#height#', f'{height - offset}')
+        template_html = template_html.replace('#width#', f'{width}')
+        template_html = template_html.replace('#src#', f'{related_path}/')
+        f.write(template_html)
+
+    rel_path = os.path.relpath(output_html_path, SAVE_DIR)
+    iframe_tag = f'<iframe src="/static/{rel_path}" height="{height}" width="100%" frameborder="0"></iframe>'
+    print(
+        f'Find html file {output_html_path}, {os.path.exists(output_html_path)}, relative HTML path is /static/{rel_path}')
+
+    return f"""
+        <div style='height: {height}; width: 100%;'>
+        {iframe_tag}
+        </div>
+    """
+
 
 from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, \
     Hunyuan3DDiTFlowMatchingPipeline
+from hy3dgen.shapegen.pipelines import export_to_trimesh
 from hy3dgen.rembg import BackgroundRemover
 
 rmbg_worker = BackgroundRemover()
@@ -75,6 +125,10 @@ if args.enable_flashvdm:
 if args.compile:
     i23d_worker.compile()
 
+floater_remove_worker = FloaterRemover()
+degenerate_face_remove_worker = DegenerateFaceRemover()
+face_reduce_worker = FaceReducer()
+
 progress=gr.Progress()
 
 @spaces.GPU(duration=60)
@@ -88,21 +142,24 @@ def gen_shape(
     target_face_num=10000,
     randomize_seed: bool = False,
 ):
-    
+    progress(0,desc="Starting")
+
     def callback(step_idx, timestep, outputs):
-        progress_value = (step_idx+1.0)/steps
+        progress_value = ((step_idx+1.0)/steps)*(0.5/1.0)
         progress(progress_value, desc=f"Mesh generating, {step_idx + 1}/{steps} steps")
 
 
     if image is None:
         raise gr.Error("Please provide either a caption or an image.")
 
+
     seed = int(randomize_seed_fn(seed, randomize_seed))
     octree_resolution = int(octree_resolution)
     save_folder = gen_save_folder()
-
+    # 先移除背景
     image = rmbg_worker(image.convert('RGB'))
 
+    # 生成模型
     generator = torch.Generator()
     generator = generator.manual_seed(int(seed))
     outputs = i23d_worker(
@@ -113,9 +170,47 @@ def gen_shape(
         octree_resolution=octree_resolution,
         num_chunks=num_chunks,
         output_type='mesh',
-        callback=callback
+        callback=callback,
+        callback_steps=1
     )
-    print(outputs)
+
+    mesh = export_to_trimesh(outputs)[0]
+
+    path = export_mesh(mesh, save_folder, textured=False)
+
+    model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH)
+
+    return model_viewer_html, path
+
+    # if args.low_vram_mode:
+    #     torch.cuda.empty_cache()
+    
+    # if path is None:
+    #     raise gr.Error('Please generate a mesh first.')
+
+    # # 简化模型
+    # print(f'exporting {path}')
+    # print(f'reduce face to {target_face_num}')
+
+    # mesh = trimesh.load(path)
+    # progress(0.5,desc="Optimizing mesh")
+
+    # mesh = floater_remove_worker(mesh)
+    # mesh = degenerate_face_remove_worker(mesh)
+    # progress(0.6,desc="Reducing mesh faces")
+    # mesh = face_reduce_worker(mesh, target_face_num)
+    # save_folder = gen_save_folder()
+
+    # file_type = "obj"
+    # path = export_mesh(mesh, save_folder, textured=False, type=file_type)
+
+    # # for preview
+    # save_folder = gen_save_folder()
+    # _ = export_mesh(mesh, save_folder, textured=False)
+    # model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH, textured=False)
+
+    # progress(1,desc="Complete")
+    # return model_viewer_html, path
 
 
 
@@ -167,6 +262,7 @@ with gr.Blocks().queue() as demo:
         with gr.Column(scale=6):
             gr.Markdown("#### Generated Mesh")
             html_export_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output')
+            path_output  = gr.Textbox(label="Mesh Path") 
                 
         with gr.Column(scale=3):
             gr.Markdown("#### Image Examples")
@@ -176,7 +272,7 @@ with gr.Blocks().queue() as demo:
     gen_button.click(
         fn=gen_shape,
         inputs=[image,num_steps,cfg_scale,seed,octree_resolution,num_chunks,target_face_num, randomize_seed], 
-        outputs=[html_export_mesh]
+        outputs=[html_export_mesh, path_output]
     )    
 
 demo.launch()