Add simple mode

app.py

@@ -59,10 +59,16 @@ def main():
 
     with gr.Blocks(css='style.css') as demo:
         gr.Markdown(TITLE)
-        gr.Markdown(DESCRIPTION)
 
         with gr.Tabs():
-            with gr.TabItem('App'):
+            with gr.TabItem('Simple Mode'):
+                run_button_simple = gr.Button('Generate')
+                result_simple = gr.Image(show_label=False,
+                                         elem_id='result-grid')
+
+            with gr.TabItem('Advanced Mode'):
+                gr.Markdown(DESCRIPTION)
+
                 with gr.Row():
                     with gr.Column():
                         with gr.Group():
@@ -87,7 +93,7 @@ def main():
                                              label='Seed')
                             run_button = gr.Button('Run')
                     with gr.Column():
-                        result = gr.Image(label='Result', elem_id='result')
+                        result = gr.Image(show_label=False, elem_id='result')
 
             with gr.TabItem('Sample Images'):
                 with gr.Row():
@@ -119,6 +125,9 @@ def main():
                                   scheduler_type,
                               ],
                               outputs=None)
+        run_button_simple.click(fn=model.run_simple,
+                                inputs=None,
+                                outputs=result_simple)
         run_button.click(fn=model.run,
                          inputs=[
                              model_name,

model.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import logging
 import os
+import random
 import sys
 
+import numpy as np
 import PIL.Image
 import torch
 from diffusers import (DDIMPipeline, DDIMScheduler, DDPMPipeline,
@@ -37,6 +39,7 @@ class Model:
         self.scheduler_type = 'DDIM'
         self.pipeline = self._load_pipeline(self.model_name,
                                             self.scheduler_type)
+        self.rng = random.Random()
 
     def _load_pipeline(self, model_name: str,
                        scheduler_type: str) -> DiffusionPipeline:
@@ -78,18 +81,21 @@ class Model:
         for name in self.MODEL_NAMES:
             self._load_pipeline(name, 'DDPM')
 
-    def generate(self, seed: int, num_steps: int) -> PIL.Image.Image:
+    def generate(self,
+                 seed: int,
+                 num_steps: int,
+                 num_images: int = 1) -> list[PIL.Image.Image]:
         logger.info('--- generate ---')
         logger.info(f'{seed=}, {num_steps=}')
 
         torch.manual_seed(seed)
         if self.scheduler_type == 'DDPM':
-            res = self.pipeline(batch_size=1,
-                                torch_device=self.device)['sample'][0]
+            res = self.pipeline(batch_size=num_images,
+                                torch_device=self.device)['sample']
         elif self.scheduler_type in ['DDIM', 'PNDM']:
-            res = self.pipeline(batch_size=1,
+            res = self.pipeline(batch_size=num_images,
                                 torch_device=self.device,
-                                num_inference_steps=num_steps)['sample'][0]
+                                num_inference_steps=num_steps)['sample']
         else:
             raise ValueError
 
@@ -106,4 +112,23 @@ class Model:
         self.set_pipeline(model_name, scheduler_type)
         if scheduler_type == 'PNDM':
             num_steps = max(4, min(num_steps, 100))
-        return self.generate(seed, num_steps)
+        return self.generate(seed, num_steps)[0]
+
+    @staticmethod
+    def to_grid(images: list[PIL.Image.Image],
+                ncols: int = 2) -> PIL.Image.Image:
+        images = [np.asarray(image) for image in images]
+        nrows = (len(images) + ncols - 1) // ncols
+        h, w = images[0].shape[:2]
+        d = nrows * ncols - len(images)
+        if d > 0:
+            images += [np.full((h, w, 3), 255, dtype=np.uint8)] * d
+        grid = np.asarray(images).reshape(2, 2, h, w, 3).transpose(
+            0, 2, 1, 3, 4).reshape(2 * h, 2 * w, 3)
+        return PIL.Image.fromarray(grid)
+
+    def run_simple(self) -> PIL.Image.Image:
+        self.set_pipeline(self.MODEL_NAMES[0], 'DDIM')
+        seed = self.rng.randint(0, 100000)
+        images = self.generate(seed, num_steps=10, num_images=4)
+        return self.to_grid(images, 2)

style.css

@@ -1,6 +1,10 @@
 h1 {
   text-align: center;
 }
+div#result-grid {
+  max-width: 600px;
+  max-height: 600px;
+}
 div#result {
   max-width: 400px;
   max-height: 400px;