restructure code

app.py

@@ -154,61 +154,66 @@ def preprocess(input_image, do_remove_background):
     return input_image
 
 
-def pipeline_callback(output_queue: SimpleQueue, pipe: Any, step_index: int, timestep: float, callback_kwargs: dict[str, Any]) -> dict[str, Any]:
+def pipeline_callback(log_queue: SimpleQueue, pipe: Any, step_index: int, timestep: float, callback_kwargs: dict[str, Any]) -> dict[str, Any]:
     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]
 
-    output_queue.put(("log", "mvs/image", rr.Image(image)))
-    output_queue.put(("log", "mvs/latents", rr.Tensor(latents.squeeze())))
+    log_queue.put(("mvs", rr.Image(image)))
+    log_queue.put(("latents", rr.Tensor(latents.squeeze())))
 
     return callback_kwargs
 
-def generate_mvs(input_image, sample_steps, sample_seed):
+def generate_mvs(log_queue, input_image, sample_steps, sample_seed):
 
     seed_everything(sample_seed)
 
-
-    def thread_target(output_queue, input_image, sample_steps):
-        z123_image = pipeline(
-            input_image,
-            num_inference_steps=sample_steps,
-            callback_on_step_end=lambda *args, **kwargs: pipeline_callback(output_queue, *args, **kwargs),
-        ).images[0]
-        output_queue.put(("z123_image", z123_image))
-
-
-    output_queue = SimpleQueue()
-    z123_thread = threading.Thread(
-        target=thread_target,
-        args=
-        [
-            output_queue,
-            input_image,
-            sample_steps,
-        ]
-    )
-    z123_thread.start()
-
-    while True:
-        msg = output_queue.get()
-        yield msg
-        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])
-    handle.start()
-    while True:
-        msg = output_queue.get()
-        yield msg
-        if msg[0] == "mesh":
-            break
-    handle.join()
-
-def _make3d(output_queue: SimpleQueue, images: Image.Image):
+    return pipeline(
+        input_image,
+        num_inference_steps=sample_steps,
+        callback_on_step_end=lambda *args, **kwargs: pipeline_callback(log_queue, *args, **kwargs),
+    ).images[0]
+
+    # def thread_target(output_queue, input_image, sample_steps):
+    #     z123_image = pipeline(
+    #         input_image,
+    #         num_inference_steps=sample_steps,
+    #         callback_on_step_end=lambda *args, **kwargs: pipeline_callback(output_queue, *args, **kwargs),
+    #     ).images[0]
+    #     log_queue.put(("z123_image", z123_image))
+
+
+    # output_queue = SimpleQueue()
+    # z123_thread = threading.Thread(
+    #     target=thread_target,
+    #     args=
+    #     [
+    #         output_queue,
+    #         input_image,
+    #         sample_steps,
+    #     ]
+    # )
+    # z123_thread.start()
+
+    # while True:
+    #     msg = output_queue.get()
+    #     yield msg
+    #     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])
+#     handle.start()
+#     while True:
+#         msg = output_queue.get()
+#         yield msg
+#         if msg[0] == "mesh":
+#             break
+#     handle.join()
+
+def make3d(log_queue, images: Image.Image):
     global model
     if IS_FLEXICUBES:
         model.init_flexicubes_geometry(device, use_renderer=False)
@@ -238,9 +243,8 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
 
         vertices, faces, vertex_colors = mesh_out
 
-        output_queue.put(
+        log_queue.put(
             (
-                "log", 
                 "mesh", 
                 rr.Mesh3D(
                     vertex_positions=vertices, 
@@ -249,7 +253,8 @@ def _make3d(output_queue: SimpleQueue, images: Image.Image):
                 ),
             )
         )
-    output_queue.put(("mesh", mesh_out))
+    
+    return mesh_out
 
 def generate_blueprint() -> rrb.Blueprint:
     return rrb.Blueprint(
@@ -258,49 +263,56 @@ def generate_blueprint() -> rrb.Blueprint:
             rrb.Grid(
                 rrb.Spatial2DView(origin="z123image"),
                 rrb.Spatial2DView(origin="preprocessed_image"),
-                rrb.Spatial2DView(origin="mvs/image"),
-                rrb.TensorView(origin="mvs/latents", ),
+                rrb.Spatial2DView(origin="mvs"),
+                rrb.TensorView(origin="latents", ),
             ),
             column_shares=[1, 1],
         ),
+        
         collapse_panels=True,
     )
 
+def compute(log_queue, input_image, do_remove_background, sample_steps, sample_seed):
+
+    preprocessed_image = preprocess(input_image, do_remove_background)
+
+    log_queue.put(("preprocessed_image", rr.Image(preprocessed_image)))
+    # rr.log("preprocessed_image", rr.Image(preprocessed_image))
+
+    z123_image = generate_mvs(log_queue, preprocessed_image, sample_steps, sample_seed)
+    
+    log_queue.put(("z123image", rr.Image(z123_image)))
+    # rr.log("z123image", rr.Image(z123_image))
+
+    mesh_out = make3d(log_queue, z123_image)
+
+    log_queue.put("done")
+
+    
 @spaces.GPU
 @rr.thread_local_stream("InstantMesh")
 def log_to_rr(input_image, do_remove_background, sample_steps, sample_seed):
 
+    log_queue = SimpleQueue()
+    
     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(preprocessed_image, sample_steps, sample_seed):
-        if msg[0] == "z123_image":
-            z123_image = msg[1]
+    handle = threading.Thread(target=compute, args=[log_queue, input_image, do_remove_background, sample_steps, sample_seed])
+    handle.start()
+    while True:
+        msg = log_queue.get()
+        if msg == "done":
             break
-        elif msg[0] == "log":
-            entity_path = msg[1]
-            entity = msg[2]
+        else:
+            entity_path, entity = msg
             rr.log(entity_path, entity)
             yield stream.read()
-
-    rr.log("z123image", rr.Image(z123_image))
-    yield stream.read()
-
-    for msg in make3d(z123_image):
-        if msg[0] == "log":
-            rr.log(msg[1], msg[2])
-            yield stream.read()
-        if msg[0] == "mesh":
-            mesh = msg[1]
-
+    handle.join()
+    
     # return mesh
 
 _HEADER_ = '''