debug printing for make3d

app.py

@@ -12,11 +12,13 @@ from omegaconf import OmegaConf
 from einops import rearrange, repeat
 from tqdm import tqdm
 import threading
+from queue import SimpleQueue
 from typing import Any
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 import rerun as rr
 from gradio_rerun import Rerun
 
+import src
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
     FOV_to_intrinsics, 
@@ -25,6 +27,7 @@ from src.utils.camera_util import (
 )
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground, images_to_video
+from src.models.lrm_mesh import InstantMesh
 
 import tempfile
 from functools import partial
@@ -126,7 +129,7 @@ print(f'type(pipeline)={type(pipeline)}')
 # load reconstruction model
 print('Loading reconstruction model ...')
 model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="instant_mesh_large.ckpt", repo_type="model")
-model = instantiate_from_config(model_config)
+model: InstantMesh = instantiate_from_config(model_config)
 state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
 state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.') and 'source_camera' not in k}
 model.load_state_dict(state_dict, strict=True)
@@ -152,29 +155,30 @@ def preprocess(input_image, do_remove_background):
     return input_image
 
 
-def pipeline_callback(pipe: Any, step_index: int, timestep: float, callback_kwargs: dict[str, Any]) -> dict[str, Any]:
+def pipeline_callback(output_queue: SimpleQueue, pipe: Any, step_index: int, timestep: float, callback_kwargs: dict[str, Any]) -> dict[str, Any]:
     rr.set_time_sequence("iteration", step_index)
     rr.set_time_seconds("timestep", timestep)
     latents = callback_kwargs["latents"]
     image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]  # type: ignore[attr-defined]
     image = pipe.image_processor.postprocess(image, output_type="np").squeeze()  # type: ignore[attr-defined]
 
-    rr.log("output", rr.Image(image))
-    rr.log("latents", rr.Tensor(latents.squeeze()))
+    output_queue.put(("log", "mvs/image", rr.Image(image)))
+    output_queue.put(("log", "mvs/latents", rr.Tensor(latents.squeeze())))
+
     return callback_kwargs
 
 @spaces.GPU
-def generate_mvs(input_image, sample_steps, sample_seed):
-
-    print(threading.get_ident())
+def generate_mvs(input_image, sample_steps, sample_seed, output_queue: SimpleQueue):
 
     seed_everything(sample_seed)
 
-    return pipeline(
+    z123_image = pipeline(
         input_image,
         num_inference_steps=sample_steps,
-        callback_on_step_end=pipeline_callback,
-    )
+        callback_on_step_end=lambda *args, **kwargs: pipeline_callback(output_queue, *args, **kwargs),
+    ).images[0]
+
+    output_queue.put(("z123_image", z123_image))
 
     # sampling
     # z123_image = pipeline(
@@ -190,10 +194,9 @@ def generate_mvs(input_image, sample_steps, sample_seed):
 
     # return z123_image, show_image
 
-
 @spaces.GPU
-def make3d(images):
-
+def make3d(output_queue: SimpleQueue, images: Image.Image):
+    print(f'type(images)={type(images)}')
     global model
     if IS_FLEXICUBES:
         model.init_flexicubes_geometry(device, use_renderer=False)
@@ -205,9 +208,12 @@ def make3d(images):
 
     input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).to(device)
     render_cameras = get_render_cameras(batch_size=1, radius=2.5, is_flexicubes=IS_FLEXICUBES).to(device)
+    print(f'type(input_cameras)={type(input_cameras)}')
 
     images = images.unsqueeze(0).to(device)
     images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1)
+    print(f'type(images)={type(images)}')
+
 
     mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name
     print(mesh_fpath)
@@ -219,26 +225,31 @@ def make3d(images):
     with torch.no_grad():
         # get triplane
         planes = model.forward_planes(images, input_cameras)
+        print(f'type(planes)={type(planes)}')
 
         # # get video
-        # chunk_size = 20 if IS_FLEXICUBES else 1
-        # render_size = 384
+        chunk_size = 20 if IS_FLEXICUBES else 1
+        render_size = 384
         
         # frames = []
-        # for i in tqdm(range(0, render_cameras.shape[1], chunk_size)):
-        #     if IS_FLEXICUBES:
-        #         frame = model.forward_geometry(
-        #             planes,
-        #             render_cameras[:, i:i+chunk_size],
-        #             render_size=render_size,
-        #         )['img']
-        #     else:
-        #         frame = model.synthesizer(
-        #             planes,
-        #             cameras=render_cameras[:, i:i+chunk_size],
-        #             render_size=render_size,
-        #         )['images_rgb']
-        #     frames.append(frame)
+        for i in tqdm(range(0, render_cameras.shape[1], chunk_size)):
+            if IS_FLEXICUBES:
+                frame = model.forward_geometry(
+                    planes,
+                    render_cameras[:, i:i+chunk_size],
+                    render_size=render_size,
+                )['img']
+            else:
+                frame = model.synthesizer(
+                    planes,
+                    cameras=render_cameras[:, i:i+chunk_size],
+                    render_size=render_size,
+                )['images_rgb']
+
+            print(f'type(framee)={type(frame)}')
+            output_queue.put(("log", "3dvideo", rr.Image(frame)))
+            # frames.append(frame)
+
         # frames = torch.cat(frames, dim=1)
 
         # images_to_video(
@@ -255,10 +266,13 @@ def make3d(images):
             use_texture_map=False,
             **infer_config,
         )
+        print(f'type(mesh_out)={type(mesh_out)}')
+
 
         vertices, faces, vertex_colors = mesh_out
         vertices = vertices[:, [1, 2, 0]]
-        
+        print(f'type(vertices)={type(vertices)}')
+
         save_glb(vertices, faces, vertex_colors, mesh_glb_fpath)
         save_obj(vertices, faces, vertex_colors, mesh_fpath)
         
@@ -266,31 +280,47 @@ def make3d(images):
 
     return mesh_fpath, mesh_glb_fpath
 
-@spaces.GPU
-def print_thread_ident_from_gpu():
-    print(threading.get_ident())
-
 @rr.thread_local_stream("InstantMesh")
 def log_to_rr(input_image, do_remove_background, sample_steps, sample_seed):
+    preprocessed_image = preprocess(input_image, do_remove_background)
+
+    stream = rr.binary_stream()
+
+    rr.log("preprocessed_image", rr.Image(preprocessed_image))
 
-    print(threading.get_ident())
-    print_thread_ident_from_gpu()
+    yield stream.read()
 
-    # preprocessed_image = preprocess(input_image, do_remove_background)
+    output_queue = SimpleQueue()
 
-    # stream = rr.binary_stream()
+    mvs_thread = threading.Thread(target=generate_mvs, args=[input_image, sample_steps, sample_seed, output_queue])
+    mvs_thread.start()
 
-    # rr.log("preprocessed_image", rr.Image(preprocessed_image))
+    while True:
+        msg = output_queue.get()
+        if msg[0] == "z123_image":
+            z123_image = msg[1]
+            break
+        elif msg[0] == "log":
+            entity_path = msg[1]
+            entity = msg[2]
+            rr.log(entity_path, entity)
+            yield stream.read()
+    
+    mvs_thread.join()
 
-    # yield stream.read()
+    rr.log("z123image", rr.Image(z123_image))
+    yield stream.read()
 
-    # z123_out = generate_mvs(input_image, sample_steps, sample_seed)
-    # print(z123_out)
-    # for image in z123_out.images:
-    #     rr.log("z123image", rr.Image(image))
-    #     yield stream.read()
-    # yield stream.read()
-    # pass
+    mesh_fpath, mesh_glb_fpath = make3d(output_queue, z123_image)
+
+    while not output_queue.empty():
+        msg = output_queue.get()
+        if msg[0] == "log":
+            entity_path = msg[1]
+            entity = msg[2]
+            rr.log(entity_path, entity)
+            yield stream.read()
+        
 
 _HEADER_ = '''
 <h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/TencentARC/InstantMesh' target='_blank'><b>InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models</b></a></h2>
@@ -343,14 +373,6 @@ with gr.Blocks() as demo:
                     type="pil",
                     elem_id="content_image",
                 )
-                processed_image = gr.Image(
-                    label="Processed Image", 
-                    image_mode="RGBA", 
-                    #width=256,
-                    #height=256,
-                    type="pil", 
-                    interactive=False
-                )
             with gr.Row():
                 with gr.Group():
                     do_remove_background = gr.Checkbox(