Add an option to show denoising process

app.py

@@ -22,7 +22,7 @@ def create_simple_demo(model: Model) -> gr.Blocks:
 
 
 def create_advanced_demo(model: Model) -> gr.Blocks:
-    def update_scheduler_type(name: str) -> dict:
+    def update_num_steps(name: str) -> dict:
         visible = name != 'DDPM'
         if name == 'PNDM':
             minimum = 4
@@ -35,6 +35,9 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                                 maximum=maximum,
                                 value=20)
 
+    def show_denoising_changed(selected: bool) -> dict:
+        return gr.Video.update(visible=selected)
+
     with gr.Blocks() as demo:
         gr.Markdown(DESCRIPTION)
 
@@ -60,9 +63,14 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                                      step=1,
                                      value=1234,
                                      label='Seed')
+                    show_denoising = gr.Checkbox(value=False,
+                                                 label='Show Denoising')
                     run_button = gr.Button('Run')
             with gr.Column():
                 result = gr.Image(show_label=False, elem_id='result')
+                result_video = gr.Video(show_label=False,
+                                        visible=False,
+                                        elem_id='result-video')
 
         model_name.change(fn=model.set_pipeline,
                           inputs=[
@@ -70,7 +78,7 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                               scheduler_type,
                           ],
                           outputs=None)
-        scheduler_type.change(fn=update_scheduler_type,
+        scheduler_type.change(fn=update_num_steps,
                               inputs=scheduler_type,
                               outputs=num_steps,
                               queue=False)
@@ -80,6 +88,9 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                                   scheduler_type,
                               ],
                               outputs=None)
+        show_denoising.change(fn=show_denoising_changed,
+                              inputs=show_denoising,
+                              outputs=result_video)
         run_button.click(fn=model.run,
                          inputs=[
                              model_name,
@@ -87,10 +98,12 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                              num_steps,
                              randomize_seed,
                              seed,
+                             show_denoising,
                          ],
                          outputs=[
                              result,
                              seed,
+                             result_video,
                          ])
     return demo
 

model.py

@@ -4,10 +4,13 @@ import logging
 import os
 import random
 import sys
+import tempfile
 
+import imageio
 import numpy as np
 import PIL.Image
 import torch
+import tqdm.auto
 from diffusers import (DDIMPipeline, DDIMScheduler, DDPMPipeline,
                        DiffusionPipeline, PNDMPipeline, PNDMScheduler)
 
@@ -101,20 +104,58 @@ class Model:
         logger.info('--- done ---')
         return res
 
-    def run(
-        self,
-        model_name: str,
-        scheduler_type: str,
-        num_steps: int,
-        randomize_seed: bool,
-        seed: int,
-    ) -> tuple[PIL.Image.Image, int]:
+    @staticmethod
+    def postprocess(sample: torch.Tensor) -> np.ndarray:
+        res = (sample / 2 + 0.5).clamp(0, 1)
+        res = (res * 255).to(torch.uint8)
+        res = res.cpu().permute(0, 2, 3, 1).numpy()
+        return res
+
+    @torch.inference_mode()
+    def generate_with_video(self, seed: int,
+                            num_steps: int) -> tuple[PIL.Image.Image, str]:
+        logger.info('--- generate_with_video ---')
+        if self.scheduler_type == 'DDPM':
+            num_steps = 1000
+            fps = 100
+        else:
+            fps = 10
+        logger.info(f'{seed=}, {num_steps=}')
+
+        model = self.pipeline.unet.to(self.device)
+        scheduler = self.pipeline.scheduler
+        scheduler.set_timesteps(num_inference_steps=num_steps)
+        input_shape = (1, model.config.in_channels, model.config.sample_size,
+                       model.config.sample_size)
+        torch.manual_seed(seed)
+
+        out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+        writer = imageio.get_writer(out_file.name, fps=fps)
+        sample = torch.randn(input_shape).to(self.device)
+        for t in tqdm.auto.tqdm(scheduler.timesteps):
+            out = model(sample, t)['sample']
+            sample = scheduler.step(out, t, sample)['prev_sample']
+            res = self.postprocess(sample)[0]
+            writer.append_data(res)
+        writer.close()
+
+        logger.info('--- done ---')
+        return res, out_file.name
+
+    def run(self, model_name: str, scheduler_type: str, num_steps: int,
+            randomize_seed: bool, seed: int, visualize_denoising: bool
+            ) -> tuple[PIL.Image.Image, int, str | None]:
         self.set_pipeline(model_name, scheduler_type)
         if scheduler_type == 'PNDM':
             num_steps = max(4, min(num_steps, 100))
         if randomize_seed:
             seed = self.rng.randint(0, 100000)
-        return self.generate(seed, num_steps)[0], seed
+
+        if not visualize_denoising:
+            return self.generate(seed, num_steps)[0], seed, None
+        else:
+            res, filename = self.generate_with_video(seed, num_steps)
+            return res, seed, filename
 
     @staticmethod
     def to_grid(images: list[PIL.Image.Image],

style.css

@@ -9,6 +9,10 @@ div#result {
   max-width: 400px;
   max-height: 400px;
 }
+div#result-video {
+  max-width: 400px;
+  max-height: 400px;
+}
 img#visitor-badge {
   display: block;
   margin: auto;