Update app.py

app.py

@@ -6,33 +6,50 @@ from torchvision import transforms
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--ckpt_path", type=str, default="./checkpoints/model/mnist.ckpt")
 parser.add_argument("--map_location", type=str, default="cpu")
 parser.add_argument("--share", action='store_true')
 args = parser.parse_args()
 
 if __name__ == "__main__":
-    model = diffusion.DiffusionModel.load_from_checkpoint(
-        args.ckpt_path, in_channels=1, map_location=args.map_location, num_classes=10
+    model_mnist = diffusion.DiffusionModel.load_from_checkpoint(
+        "./checkpoints/model/mnist.ckpt"
+    )
+    model_celeba = diffusion.DiffusionModel.load_from_checkpoint(
+        "./checkpoints/model/celebahq.ckpt"
     )
     to_pil = transforms.ToPILImage()
 
-    def reset(image):
-        image = to_pil((torch.randn(1, 32, 32)*255).type(torch.uint8))
-        return image
-
-    # def noising(image):
-    #     for i in range(100):
+    def denoise_celeb(timesteps):
+        for img in model_celeba.sampling(demo=True, mode="ddim", timesteps=timesteps, n_samples=1):
+            image = to_pil(img[0])
+            yield image
 
     def denoise(label, timesteps):
-        labels = torch.tensor([label]).to(model.device)
-        for img in model.sampling(labels=labels, demo=True, mode="ddim", timesteps=timesteps):
+        labels = torch.tensor([label]).to(model_mnist.device)
+        for img in model_mnist.sampling(labels=labels, demo=True, mode="ddim", timesteps=timesteps):
             image = to_pil(img[0])
             yield image
 
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="green")) as demo:
         gr.Markdown("# Simple Diffusion Model")
 
+        gr.Markdown("## CelebA")
+        with gr.Row():
+            with gr.Column(scale=2):
+                timesteps_celeb = gr.Radio(
+                    label='Timestep', choices=[10, 20, 50, 100, 200, 1000]
+                )
+                sample_celeb_btn = gr.Button("Sample")
+
+            output = gr.Image(
+                value=to_pil((torch.randn(3, 64, 64)*255).type(torch.uint8)),
+                scale=1,
+                image_mode="RGB",
+                type='pil',
+            )
+
+        sample_celeb_btn.click(denoise_celeb, [timesteps_celeb], outputs=output)
+
         gr.Markdown("## MNIST")
         with gr.Row():
             with gr.Column(scale=2):
@@ -42,17 +59,17 @@ if __name__ == "__main__":
                         choices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
                         value=0
                     )
-                    timesteps = gr.Radio(label='Timestep', choices=[10, 20, 50, 100, 200, 1000])
+                    timesteps = gr.Radio(
+                        label='Timestep', choices=[10, 20, 50, 100, 200, 1000]
+                    )
                 with gr.Row():
-                    sample_btn = gr.Button("Sampling")
-                    reset_btn = gr.Button("Reset")
+                    sample_mnist_btn = gr.Button("Sample")
             output = gr.Image(
                 value=to_pil((torch.randn(1, 32, 32)*255).type(torch.uint8)),
-                scale=2,
+                scale=1,
                 image_mode="L",
                 type='pil',
             )
-            sample_btn.click(denoise, [label, timesteps], outputs=output)
-            reset_btn.click(reset, [output], outputs=output)
+            sample_mnist_btn.click(denoise, [label, timesteps], outputs=output)
 
     demo.launch(share=args.share)