add training progress visualization

gradio_demo.py

@@ -32,17 +32,25 @@ example_inputs = [[
         "A DSLR photo of a wooden car, super detailed, best quality, 4K, HD.",
         "a wooden car."
 ]]
-example_outputs = [
+example_outputs_1 = [
     gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/boots.mp4'), autoplay=True),
     gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/Donut.mp4'), autoplay=True),
     gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/durian.mp4'), autoplay=True),
     gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/pillow_huskies.mp4'), autoplay=True),
     gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/wooden_car.mp4'), autoplay=True)
 ]
+example_outputs_2 = [
+    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/boots_pro.mp4'), autoplay=True),
+    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/Donut_pro.mp4'), autoplay=True),
+    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/durian_pro.mp4'), autoplay=True),
+    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/pillow_huskies_pro.mp4'), autoplay=True),
+    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/wooden_car_pro.mp4'), autoplay=True)
+]
+
 
 def main(prompt, init_prompt, negative_prompt, num_iter, CFG, seed):
     if [prompt, init_prompt] in example_inputs:
-        return example_outputs[example_inputs.index([prompt, init_prompt])]
+        return example_outputs_1[example_inputs.index([prompt, init_prompt])], example_outputs_2[example_inputs.index([prompt, init_prompt])]
     args, lp, op, pp, gcp, gp = args_parser(default_opt=os.path.join(os.path.dirname(__file__), 'configs/white_hair_ironman.yaml'))
     gp.text = prompt
     gp.negative = negative_prompt
@@ -59,16 +67,16 @@ def main(prompt, init_prompt, negative_prompt, num_iter, CFG, seed):
     if os.environ.get('QUEUE_1') != "True":
         os.environ['QUEUE_1'] = "True"
         lp.workspace = 'gradio_demo_1'
-        video_path = start_training(args, lp, op, pp, gcp, gp)
+        video_path, pro_video_path = start_training(args, lp, op, pp, gcp, gp)
         os.environ['QUEUE_1'] = "False"
     else:
         lp.workspace = 'gradio_demo_2'
-        video_path = start_training(args, lp, op, pp, gcp, gp)
-    return gr.Video(value=video_path, autoplay=True)
+        video_path, pro_video_path = start_training(args, lp, op, pp, gcp, gp)
+    return gr.Video(value=video_path, autoplay=True), gr.Video(value=pro_video_path, autoplay=True)
 
 with gr.Blocks() as demo:
     gr.Markdown("# <center>LucidDreamer: Towards High-Fidelity Text-to-3D Generation via Interval Score Matching</center>")
-    gr.Markdown("This live demo allows you to generate high-quality 3D content using text prompts.<br> \
+    gr.Markdown("This live demo allows you to generate high-quality 3D content using text prompts. The outputs are 360° rendered 3d gaussian video and training progress visualization.<br> \
                 It is based on Stable Diffusion 2.1. Please check out our <strong><a href=https://github.com/EnVision-Research/LucidDreamer>Project Page</a> / <a href=https://arxiv.org/abs/2311.11284>Paper</a> / <a href=https://github.com/EnVision-Research/LucidDreamer>Code</a></strong> if you want to learn more about our method!<br> \
                 Note that this demo is running on A10G Small, the running time might be longer than the reported 35 minutes (5000 iterations) on A100.<br> \
                 &copy; This Gradio space was developed by Haodong LI.")
@@ -78,7 +86,7 @@ with gr.Blocks() as demo:
             gr.Slider(1000, 5000, value=5000, label="Number of iterations"),
             gr.Slider(7.5, 100, value=7.5, label="CFG"),
             gr.Number(value=0, label="Seed")], 
-        outputs="playable_video",
+        outputs=["playable_video", "playable_video"],
         examples=example_inputs,
         cache_examples=True,
         concurrency_limit=2)

train.py

@@ -382,7 +382,7 @@ def training(dataset, opt, pipe, gcams, guidance_opt, testing_iterations, saving
 
     if opt.save_process:
         imageio.mimwrite(os.path.join(save_folder_proc, "video_rgb.mp4"), pro_img_frames, fps=30, quality=8)
-    return video_path
+    return video_path, os.path.join(save_folder_proc, "video_rgb.mp4")
 
 
 def prepare_output_and_logger(args):    
@@ -543,10 +543,10 @@ def start_training(args, lp, op, pp, gcp, gp):
     # Start GUI server, configure and run training
     network_gui.init(args.ip, args.port)
     torch.autograd.set_detect_anomaly(args.detect_anomaly)
-    video_path = training(lp, op, pp, gcp, gp, args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from, args.save_video)
+    video_path, pro_video_path = training(lp, op, pp, gcp, gp, args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from, args.save_video)
     # All done
     print("\nTraining complete.")
-    return video_path
+    return video_path, pro_video_path
 
 if __name__ == "__main__":
     args, lp, op, pp, gcp, gp = args_parser()