Advanced 3D settings

.gitignore

@@ -163,7 +163,7 @@ cython_debug/
 /.vs
 /src/__pycache__
 /utils/__pycache__
-/__pycache__
+**/__pycache__
 /temp_models
 /.vscode/settings.json
 **/*.pyc

app.py

@@ -613,7 +613,7 @@ def combine_images_with_lerp(input_image, output_image, alpha):
     print(f"Combining images with alpha: {alpha}")
     return lerp_imagemath(in_image, out_image, alpha)
 
-def add_border(image, mask_width, mask_height, blank_color):    
+def add_border(image, mask_width, mask_height, blank_color):
     bordered_image_output = Image.open(image).convert("RGBA")
     margin_color = detect_color_format(blank_color)
     print(f"Adding border to image with width: {mask_width}, height: {mask_height}, color: {margin_color}")
@@ -751,47 +751,54 @@ def generate_3d_asset_part1(depth_image_source, randomize_seed, seed, input_imag
     # Determine the final seed using default MAX_SEED from constants
     final_seed = np.random.randint(0, constants.MAX_SEED) if randomize_seed else seed
     # Process the image for depth estimation
-    depth_img = depth_process_image(image_path, resized_width=1536, z_scale=332)
+    depth_img = depth_process_image(image_path, resized_width=1536, z_scale=336)
     depth_img = resize_image_with_aspect_ratio(depth_img, 1536, 1536)
 
     return depth_img, image_path, output_name, final_seed
 
 @spaces.GPU(duration=150,progress=gr.Progress(track_tqdm=True))
-def generate_3d_asset_part2(depth_img, image_path, output_name, seed, req: gr.Request, progress=gr.Progress(track_tqdm=True)):
+def generate_3d_asset_part2(depth_img, image_path, output_name, seed, steps, model_resolution, video_resolution, req: gr.Request, progress=gr.Progress(track_tqdm=True)):
     # Open image using standardized defaults
     image_raw = Image.open(image_path).convert("RGB")
-    resized_image = resize_image_with_aspect_ratio(image_raw, 1536, 1536)
+    resized_image = resize_image_with_aspect_ratio(image_raw, model_resolution, model_resolution)
     depth_img = Image.open(depth_img).convert("RGBA")
     # Preprocess and run the Trellis pipeline with fixed sampler settings
-    processed_image = TRELLIS_PIPELINE.preprocess_image(resized_image, max_resolution=1536)
-    outputs = TRELLIS_PIPELINE.run(
-        processed_image,
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": 15,
-            "cfg_strength": 7.5,
-        },
-        slat_sampler_params={
-            "steps": 15,
-            "cfg_strength": 3.0,
-        },
-    )
+    try:
+        processed_image = TRELLIS_PIPELINE.preprocess_image(resized_image, max_resolution=model_resolution)
+        outputs = TRELLIS_PIPELINE.run(
+            processed_image,
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            preprocess_image=False,
+            sparse_structure_sampler_params={
+                "steps": steps,
+                "cfg_strength": 7.5,
+            },
+            slat_sampler_params={
+                "steps": steps,
+                "cfg_strength": 3.0,
+            },
+        )
 
-    # Validate the mesh
-    mesh = outputs['mesh'][0]
-    meshisdict = isinstance(mesh, dict)
-    if meshisdict:
-        vertices = mesh['vertices']
-        faces = mesh['faces']
-    else:
-        vertices = mesh.vertices
-        faces = mesh.faces
+        # Validate the mesh
+        mesh = outputs['mesh'][0]
+        meshisdict = isinstance(mesh, dict)
+        if meshisdict:
+            vertices = mesh['vertices']
+            faces = mesh['faces']
+        else:
+            vertices = mesh.vertices
+            faces = mesh.faces
 
-    print(f"Mesh vertices: {vertices.shape}, faces: {faces.shape}")
-    if faces.max() >= vertices.shape[0]:
-        raise ValueError(f"Invalid mesh: face index {faces.max()} exceeds vertex count {vertices.shape[0]}")
+        print(f"Mesh vertices: {vertices.shape}, faces: {faces.shape}")
+        if faces.max() >= vertices.shape[0]:
+            raise ValueError(f"Invalid mesh: face index {faces.max()} exceeds vertex count {vertices.shape[0]}")
+    except Exception as e:
+        gr.Warning(f"Error generating 3D asset: {e}")
+        print(f"Error generating 3D asset: {e}")
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+        return None,None, depth_img
 
     # Ensure data is on GPU and has correct type
     if not vertices.is_cuda or not faces.is_cuda:
@@ -807,9 +814,13 @@ def generate_3d_asset_part2(depth_img, image_path, output_name, seed, req: gr.Re
     user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
 
-    video = render_utils.render_video(outputs['gaussian'][0], resolution=576, num_frames=64, r=1, fov=45)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], resolution=576, num_frames=64, r=1, fov=45)['normal']
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    video = render_utils.render_video(outputs['gaussian'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['color']
+    try:
+        video_geo = render_utils.render_video(outputs['mesh'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['normal']
+        video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    except Exception as e:
+        gr.Info(f"Error rendering video: {e}")
+        print(f"Error rendering video: {e}")
     video_path = os.path.join(user_dir, f'{output_name}.mp4')
     imageio.mimsave(video_path, video, fps=8)
 
@@ -818,7 +829,9 @@ def generate_3d_asset_part2(depth_img, image_path, output_name, seed, req: gr.Re
     depth_snapshot = depth_img
 
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], output_name)
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
     return [state, video_path, depth_snapshot]
 
 
@@ -845,7 +858,9 @@ def extract_glb(
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, f'{name}.glb')
     glb.export(glb_path)
+
     torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
     return glb_path, glb_path
 
 @spaces.GPU(progress=gr.Progress(track_tqdm=True))
@@ -863,7 +878,9 @@ def extract_gaussian(state: dict, req: gr.Request, progress=gr.Progress(track_tq
     gs, _, name = unpack_state(state)
     gaussian_path = os.path.join(user_dir, f'{name}.ply')
     gs.save_ply(gaussian_path)
+
     torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
     return gaussian_path, gaussian_path
 
 
@@ -1171,16 +1188,21 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
 
     with gr.Accordion("Height Maps and 3D", open=False):
         with gr.Row():
-            with gr.Column():
-                # Use standard seed settings only
-                seed_3d = gr.Slider(0, constants.MAX_SEED, label="Seed (3D Generation)", value=0, step=1, randomize=True)
-                randomize_seed_3d = gr.Checkbox(label="Randomize Seed (3D Generation)", value=True)
-            with gr.Column():
-                depth_image_source = gr.Radio(
-                    label="Depth Image Source",
-                    choices=["Input Image", "Output Image", "Overlay Image", "Image with Margins"],
-                    value="Input Image"
-                )
+            depth_image_source = gr.Radio(
+                label="Depth Image Source",
+                choices=["Input Image", "Hexagon Grid Image", "Overlay Image", "Image with Margins"],
+                value="Input Image"
+            )
+        with gr.Accordion("Advanced 3D Generation Settings", open=False):
+            with gr.Row():
+                with gr.Column():
+                    # Use standard seed settings only
+                    seed_3d = gr.Slider(0, constants.MAX_SEED, label="Seed (3D Generation)", value=0, step=1, randomize=True)
+                    randomize_seed_3d = gr.Checkbox(label="Randomize Seed (3D Generation)", value=True)
+                with gr.Column():
+                    steps = gr.Slider(6, 36, value=12, step=1, label="Image Sampling Steps", interactive=True)
+                    video_resolution = gr.Slider(384, 768, value=480, step=32, label="Video Resolution (*danger*)", interactive=True)
+                    model_resolution = gr.Slider(512, 2304, value=1024, step=64, label="3D Model Resolution", interactive=True)
         with gr.Row():
             generate_3d_asset_button = gr.Button("Generate 3D Asset", elem_classes="solid", variant="secondary")
         with gr.Row():
@@ -1196,9 +1218,16 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
                 extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                 extract_gaussian_btn = gr.Button("Extract Gaussian", interactive=False)
             with gr.Row():
-                model_output = gr.Model3D(label="Extracted 3D Model", clear_color=[1.0, 1.0, 1.0, 1.0], 
+                with gr.Column(scale=2):
+                    model_output = gr.Model3D(label="Extracted 3D Model", clear_color=[1.0, 1.0, 1.0, 1.0], 
                                             elem_classes="centered solid imgcontainer", interactive=True)
-                model_file = gr.File(label="3D GLTF", elem_classes="solid small centered")
+                with gr.Column(scale=1):
+                    glb_file = gr.File(label="3D GLTF", elem_classes="solid small centered", height=250)
+                    gaussian_file = gr.File(label="Gaussian", elem_classes="solid small centered", height=250)
+                    gr.Markdown("""
+                    ### Files over 10 MB may not display in the 3D model viewer
+                    """, elem_id="file_size_info", elem_classes="intro" )
+
     is_multiimage = gr.State(False)
     output_buf = gr.State()
     ddd_image_path = gr.State("./images/images/Beeuty-1.png")
@@ -1320,7 +1349,7 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
         scroll_to_output=True
     ).then(
         fn=generate_3d_asset_part2,
-        inputs=[depth_output, ddd_image_path, ddd_file_name, seed_3d ],
+        inputs=[depth_output, ddd_image_path, ddd_file_name, seed_3d, steps, model_resolution, video_resolution ],
         outputs=[output_buf, video_output, depth_output],
         scroll_to_output=True
     ).then(
@@ -1332,19 +1361,19 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
     extract_glb_btn.click(
         fn=extract_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, model_file]
+        outputs=[model_output, glb_file]
     ).then(
         lambda: gr.Button(interactive=True),
-        outputs=[model_file]
+        outputs=[glb_file]
     )
 
     extract_gaussian_btn.click(
         fn=extract_gaussian,
         inputs=[output_buf],
-        outputs=[model_output, model_file]
+        outputs=[model_output, gaussian_file]
     ).then(
         lambda: gr.Button(interactive=True),
-        outputs=[model_file]
+        outputs=[gaussian_file]
     )
 
 if __name__ == "__main__":