add blueprint, fix bug where it wouldn't use the preprocessed image

app.py

@@ -16,6 +16,7 @@ from queue import SimpleQueue
 from typing import Any
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 import rerun as rr
+import rerun.blueprint as rrb
 from gradio_rerun import Rerun
 
 import src
@@ -156,15 +157,10 @@ def preprocess(input_image, do_remove_background):
 
 
 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("mvs/image", rr.Image(image))
-    # rr.log("mvs/latents", rr.Tensor(latents.squeeze()))
-
     output_queue.put(("log", "mvs/image", rr.Image(image)))
     output_queue.put(("log", "mvs/latents", rr.Tensor(latents.squeeze())))
 
@@ -174,6 +170,7 @@ def generate_mvs(input_image, sample_steps, sample_seed):
 
     seed_everything(sample_seed)
 
+
     def thread_target(output_queue, input_image, sample_steps):
         z123_image = pipeline(
             input_image,
@@ -182,6 +179,7 @@ def generate_mvs(input_image, sample_steps, sample_seed):
         ).images[0]
         output_queue.put(("z123_image", z123_image))
 
+
     output_queue = SimpleQueue()
     z123_thread = threading.Thread(
         target=thread_target,
@@ -197,11 +195,10 @@ def generate_mvs(input_image, sample_steps, sample_seed):
     while True:
         msg = output_queue.get()
         yield msg
-        if msg[0] == "z123_image":
-            break    
+        if msg[0] == "z123_image":            
+            break
     z123_thread.join()
 
-
 def make3d(images: Image.Image):
     output_queue = SimpleQueue()
     handle = threading.Thread(target=_make3d, args=[output_queue, images])
@@ -214,7 +211,6 @@ def make3d(images: Image.Image):
     handle.join()
 
 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)
@@ -226,15 +222,11 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
 
     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)
     mesh_basename = os.path.basename(mesh_fpath).split('.')[0]
     mesh_dirname = os.path.dirname(mesh_fpath)
     video_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.mp4")
@@ -243,7 +235,6 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
     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
@@ -283,7 +274,6 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
             use_texture_map=False,
             **infer_config,
         )
-        print(f'type(mesh_out)={type(mesh_out)}')
 
         vertices, faces, vertex_colors = mesh_out
 
@@ -301,10 +291,6 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
 
         vertices = vertices[:, [1, 2, 0]]
 
-        print(f'type(vertices)={type(vertices)}')
-        print(f'type(faces)={type(faces)}')
-        print(f'type(vertex_colors)={type(vertex_colors)}')
-
         save_glb(vertices, faces, vertex_colors, mesh_glb_fpath)
         save_obj(vertices, faces, vertex_colors, mesh_fpath)
 
@@ -312,19 +298,37 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
     
     output_queue.put(("mesh", mesh_out))
 
+def generate_blueprint() -> rrb.Blueprint:
+    return rrb.Blueprint(
+        rrb.Horizontal(
+            rrb.Spatial3DView(origin="mesh"),
+            rrb.Grid(
+                rrb.Spatial2DView(origin="z123image"),
+                rrb.Spatial2DView(origin="preprocessed_image"),
+                rrb.Spatial2DView(origin="mvs/image"),
+                rrb.TensorView(origin="mvs/latents"),   
+            ),
+            column_shares=[1, 1],
+        ),
+        collapse_panels=True,
+    )
 
 @spaces.GPU
 @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()
 
+    blueprint = generate_blueprint()
+    rr.send_blueprint(blueprint)
+    yield stream.read()
+
+    preprocessed_image = preprocess(input_image, do_remove_background)
     rr.log("preprocessed_image", rr.Image(preprocessed_image))
 
     yield stream.read()
 
-    for msg in generate_mvs(input_image, sample_steps, sample_seed):
+    for msg in generate_mvs(preprocessed_image, sample_steps, sample_seed):
         if msg[0] == "z123_image":
             z123_image = msg[1]
             break
@@ -426,7 +430,7 @@ with gr.Blocks() as demo:
                     examples_per_page=16
                 )
 
-        with gr.Column(scale=4):
+        with gr.Column(scale=2):
             
             viewer = Rerun(streaming=True, height=800)