Upload gradio_app.py with huggingface_hub

gradio_app.py

@@ -33,7 +33,6 @@ import numpy as np
 import PIL
 import torch
 from pathlib import Path
-import threading
 from queue import SimpleQueue
 import trimesh
 import subprocess
@@ -96,8 +95,31 @@ def svd_render_threaded(
         log_queue.put(frames)
 
 
+def svd_render(
+    image_o: PIL.Image.Image,
+    masks: Float64[torch.Tensor, "b 72 128"],
+    cond_image: PIL.Image.Image,
+    lambda_ts: Float64[torch.Tensor, "n b"],
+    num_denoise_iters: Literal[2, 25, 50, 100],
+    weight_clamp: float,
+    log_queue: SimpleQueue | None = None,
+):
+    frames: list[PIL.Image.Image] = SVD_PIPE(
+        [image_o],
+        log_queue=None,
+        temp_cond=cond_image,
+        mask=masks,
+        lambda_ts=lambda_ts,
+        weight_clamp=weight_clamp,
+        num_frames=25,
+        decode_chunk_size=8,
+        num_inference_steps=num_denoise_iters,
+    ).frames[0]
+    return frames
+
+
 if IN_SPACES:
-    svd_render_threaded = spaces.GPU(svd_render_threaded)
+    svd_render = spaces.GPU(svd_render)
     image_to_depth = spaces.GPU(image_to_depth)
 
 
@@ -207,44 +229,53 @@ def gradio_warped_image(
     progress(0.15, desc="Starting diffusion")
 
     # to allow logging from a separate thread
-    log_queue: SimpleQueue = SimpleQueue()
-    handle = threading.Thread(
-        target=svd_render_threaded,
-        kwargs={
-            "image_o": rgb_resized,
-            "masks": masks,
-            "cond_image": cond_image,
-            "lambda_ts": lambda_ts,
-            "num_denoise_iters": num_denoise_iters,
-            "weight_clamp": 0.2,
-            "log_queue": None,
-        },
+    # log_queue: SimpleQueue = SimpleQueue()
+    # handle = threading.Thread(
+    #     target=svd_render_threaded,
+    #     kwargs={
+    #         "image_o": rgb_resized,
+    #         "masks": masks,
+    #         "cond_image": cond_image,
+    #         "lambda_ts": lambda_ts,
+    #         "num_denoise_iters": num_denoise_iters,
+    #         "weight_clamp": 0.2,
+    #         "log_queue": None,
+    #     },
+    # )
+
+    # handle.start()
+    # i = 0
+    # while True:
+    #     msg = log_queue.get()
+    #     match msg:
+    #         case frames if all(isinstance(frame, PIL.Image.Image) for frame in frames):
+    #             break
+    #         case entity_path, entity, times:
+    #             i += 1
+    #             rr.reset_time()
+    #             for timeline, time in times:
+    #                 if isinstance(time, int):
+    #                     rr.set_time_sequence(timeline, time)
+    #                 else:
+    #                     rr.set_time_seconds(timeline, time)
+    #             static = False
+    #             if entity_path == "diffusion_step":
+    #                 static = True
+    #             rr.log(entity_path, entity, static=static)
+    #             yield stream.read(), None, [], f"{i} out of {num_denoise_iters}"
+    #         case _:
+    #             assert False
+    # handle.join()
+    frames = svd_render(
+        image_o=rgb_resized,
+        masks=masks,
+        cond_image=cond_image,
+        lambda_ts=lambda_ts,
+        num_denoise_iters=num_denoise_iters,
+        weight_clamp=0.2,
+        log_queue=None,
     )
 
-    handle.start()
-    i = 0
-    while True:
-        msg = log_queue.get()
-        match msg:
-            case frames if all(isinstance(frame, PIL.Image.Image) for frame in frames):
-                break
-            case entity_path, entity, times:
-                i += 1
-                rr.reset_time()
-                for timeline, time in times:
-                    if isinstance(time, int):
-                        rr.set_time_sequence(timeline, time)
-                    else:
-                        rr.set_time_seconds(timeline, time)
-                static = False
-                if entity_path == "diffusion_step":
-                    static = True
-                rr.log(entity_path, entity, static=static)
-                yield stream.read(), None, [], f"{i} out of {num_denoise_iters}"
-            case _:
-                assert False
-    handle.join()
-
     # all frames but the first one
     frame: np.ndarray
     for frame_id, (frame, cam_pararms) in enumerate(zip(frames, camera_list)):