Gradio Working!

app.py

@@ -30,17 +30,17 @@ def vis_routine(y, depth):
 
 with gr.Blocks(css=".gradio-container {max-width: 512px; margin: auto;}") as demo:
     # title
-    gr.Markdown('[Score Jacobian Chaining](https://github.com/pals-ttic/sjc) Lifting Pretrained 2D Diffusion Models for 3D Generation')
+    gr.Markdown('[Score Jacobian Chaining](https://github.com/pals-ttic/sjc): Lifting Pretrained 2D Diffusion Models for 3D Generation')
 
     # inputs
     prompt = gr.Textbox(label="Prompt", max_lines=1, value="A high quality photo of a delicious burger")
-    iters = gr.Slider(label="Iters", minimum=1000, maximum=20000, value=10000, step=100)
+    iters = gr.Slider(label="Iters", minimum=100, maximum=20000, value=10000, step=100)
     seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
     button = gr.Button('Generate')
 
     # outputs
     image = gr.Image(label="image", visible=True)
-    depth = gr.Image(label="depth", visible=True)
+    # depth = gr.Image(label="depth", visible=True)
     video = gr.Video(label="video", visible=False)
     logs = gr.Textbox(label="logging")
 
@@ -141,78 +141,64 @@ with gr.Blocks(css=".gradio-container {max-width: 512px; margin: auto;}") as dem
                         if isinstance(model, StableDiffusion):
                             y = model.decode(y)
                         pane, img, depth = vis_routine(y, depth)
+                        yield {
+                            image: gr.update(value=img, visible=True),
+                            video: gr.update(visible=False),
+                            logs: f"Steps: {i}/{n_steps}: \n" + str(tsr_stats(y)),
+                        }
 
                 # TODO: Output pane, img and depth to Gradio
 
                 pbar.update()
                 pbar.set_description(p)
 
-                yield {
-                    image: gr.update(value=img, visible=True),
-                    depth: gr.update(value=depth, visible=True),
-                    video: gr.update(visible=False),
-                    logs: str(tsr_stats(y)),
-                }
-
             # TODO: Save Checkpoint
-            ckpt = vox.state_dict()
-            H, W = poser.H, poser.W
-            vox.eval()
-            K, poses = poser.sample_test(100)
+            with torch.no_grad():
+                ckpt = vox.state_dict()
+                H, W = poser.H, poser.W
+                vox.eval()
+                K, poses = poser.sample_test(100)
 
-            aabb = vox.aabb.T.cpu().numpy()
-            vox = vox.to(device_glb)
+                aabb = vox.aabb.T.cpu().numpy()
+                vox = vox.to(device_glb)
 
-            num_imgs = len(poses)
+                num_imgs = len(poses)
 
-            for i in (pbar := tqdm(range(num_imgs))):
+                all_images = []
 
-                pose = poses[i]
-                y, depth = render_one_view(vox, aabb, H, W, K, pose)
-                if isinstance(model, StableDiffusion):
-                    y = model.decode(y)
-                pane, img, depth = vis_routine(y, depth)
+                for i in (pbar := tqdm(range(num_imgs))):
 
-                # Save img to output
-                img.save(f"output/{i}.png")
+                    pose = poses[i]
+                    y, depth = render_one_view(vox, aabb, H, W, K, pose)
+                    if isinstance(model, StableDiffusion):
+                        y = model.decode(y)
+                    pane, img, depth = vis_routine(y, depth)
 
-                yield {
-                    image: gr.update(value=img, visible=True),
-                    depth: gr.update(value=depth, visible=True),
-                    video: gr.update(visible=False),
-                    logs: str(tsr_stats(y)),
-                }
-
-            output_video = "view_seq.mp4"
+                    # Save img to output
+                    all_images.append(img)
 
-            def export_movie(seqs, fname, fps=30):
-                fname = Path(fname)
-                if fname.suffix == "":
-                    fname = fname.with_suffix(".mp4")
-                writer = imageio.get_writer(fname, fps=fps)
-                for img in seqs:
-                    writer.append_data(img)
-                writer.close()
+                    yield {
+                        image: gr.update(value=img, visible=True),
+                        video: gr.update(visible=False),
+                        logs: str(tsr_stats(y)),
+                    }
 
-            def stitch_vis(save_fn, img_fnames, fps=10):
-                figs = [imageio.imread(fn) for fn in img_fnames]
-                export_movie(figs, save_fn, fps)
+                output_video = "/tmp/tmp.mp4"
 
-            stitch_vis(output_video, [f"output/{i}.png" for i in range(num_imgs)])
+                imageio.mimwrite(output_video, all_images, quality=8, fps=10)
 
-            end_t = time.time()
+                end_t = time.time()
 
-            yield {
-                image: gr.update(value=img, visible=False),
-                depth: gr.update(value=depth, visible=False),
-                video: gr.update(value=output_video, visible=True),
-                logs: f"Generation Finished in {(end_t - start_t)/ 60:.4f} minutes!",
-            }
+                yield {
+                    image: gr.update(value=img, visible=False),
+                    video: gr.update(value=output_video, visible=True),
+                    logs: f"Generation Finished in {(end_t - start_t)/ 60:.4f} minutes!",
+                }
 
     button.click(
             submit,
             [prompt, iters, seed],
-            [image, depth, video, logs]
+            [image, video, logs]
     )
 
 # concurrency_count: only allow ONE running progress, else GPU will OOM.