fix pickle error

app.py

@@ -168,31 +168,37 @@ def pipeline_callback(output_queue: SimpleQueue, pipe: Any, step_index: int, tim
     return callback_kwargs
 
 @spaces.GPU
-def generate_mvs(output_queue: SimpleQueue, input_image, sample_steps, sample_seed):
+def generate_mvs(input_image, sample_steps, sample_seed):
 
     seed_everything(sample_seed)
 
-    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]
+    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.put(("z123_image", z123_image))
-
-    # sampling
-    # z123_image = pipeline(
-    #     input_image, 
-    #     num_inference_steps=sample_steps
-    # ).images[0]
+    output_queue = SimpleQueue()
+    z123_thread = threading.Thread(
+        target=thread_target,
+        args=
+        [
+            output_queue,
+            input_image,
+            sample_steps,
+        ]
+    )
+    z123_thread.start()
 
-    # show_image = np.asarray(z123_image, dtype=np.uint8)
-    # show_image = torch.from_numpy(show_image)     # (960, 640, 3)
-    # show_image = rearrange(show_image, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
-    # show_image = rearrange(show_image, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
-    # show_image = Image.fromarray(show_image.numpy())
+    while True:
+        msg = output_queue.get()
+        yield msg
+        if msg[0] == "z123_image":
+            break    
+    z123_thread.join()
 
-    # return z123_image, show_image
 
 @spaces.GPU
 def make3d(output_queue: SimpleQueue, images: Image.Image):
@@ -292,13 +298,25 @@ def log_to_rr(input_image, do_remove_background, sample_steps, sample_seed):
 
     yield stream.read()
 
-    output_queue = SimpleQueue()
+    # output_queue = SimpleQueue()
 
-    mvs_thread = threading.Thread(target=generate_mvs, args=[output_queue, input_image, sample_steps, sample_seed])
-    mvs_thread.start()
+    # mvs_thread = threading.Thread(target=generate_mvs, args=[output_queue, input_image, sample_steps, sample_seed])
+    # mvs_thread.start()
 
-    while True:
-        msg = output_queue.get()
+    # 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()
+
+    for msg in generate_mvs(input_image, sample_steps, sample_seed):
         if msg[0] == "z123_image":
             z123_image = msg[1]
             break
@@ -307,8 +325,6 @@ def log_to_rr(input_image, do_remove_background, sample_steps, sample_seed):
             entity = msg[2]
             rr.log(entity_path, entity)
             yield stream.read()
-    
-    mvs_thread.join()
 
     rr.log("z123image", rr.Image(z123_image))
     yield stream.read()