Split the demo code

app.py

@@ -55,64 +55,41 @@ def update_scheduler_type(name: str) -> dict:
                             value=20)
 
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--device', type=str, default='cpu')
-    args = parser.parse_args()
-
-    model = Model(args.device)
-
-    with gr.Blocks(css='style.css') as demo:
-        gr.Markdown(TITLE)
-
-        with gr.Tabs():
-            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():
-                            model_name = gr.Dropdown(
-                                model.MODEL_NAMES,
-                                value=model.MODEL_NAMES[0],
-                                label='Model',
-                                interactive=False)
-                            scheduler_type = gr.Radio(
-                                choices=['DDPM', 'DDIM', 'PNDM'],
-                                value='DDIM',
-                                label='Scheduler')
-                            num_steps = gr.Slider(1,
-                                                  200,
-                                                  step=1,
-                                                  value=20,
-                                                  label='Number of Steps')
-                            seed = gr.Slider(0,
-                                             100000,
-                                             step=1,
-                                             value=1234,
-                                             label='Seed')
-                            run_button = gr.Button('Run')
-                    with gr.Column():
-                        result = gr.Image(show_label=False, elem_id='result')
-
-            with gr.TabItem('Sample Images'):
-                with gr.Row():
-                    model_name2 = gr.Dropdown([
-                        'ddpm-128-exp000 (DDPM)',
-                        'ddpm-128-exp000 (DDIM, 20 steps)',
-                    ],
-                                              value='ddpm-128-exp000 (DDPM)',
-                                              label='Model')
-                with gr.Row():
-                    text = get_sample_image_markdown(model_name2.value)
-                    sample_images = gr.Markdown(text)
-
-        gr.Markdown(FOOTER)
+def create_simple_demo(model: Model) -> gr.Blocks:
+    with gr.Blocks() as demo:
+        run_button = gr.Button('Generate')
+        result = gr.Image(show_label=False, elem_id='result-grid')
+        run_button.click(fn=model.run_simple, inputs=None, outputs=result)
+    return demo
+
+
+def create_advanced_demo(model: Model) -> gr.Blocks:
+    with gr.Blocks() as demo:
+        gr.Markdown(DESCRIPTION)
+
+        with gr.Row():
+            with gr.Column():
+                with gr.Group():
+                    model_name = gr.Dropdown(model.MODEL_NAMES,
+                                             value=model.MODEL_NAMES[0],
+                                             label='Model',
+                                             interactive=False)
+                    scheduler_type = gr.Radio(choices=['DDPM', 'DDIM', 'PNDM'],
+                                              value='DDIM',
+                                              label='Scheduler')
+                    num_steps = gr.Slider(1,
+                                          200,
+                                          step=1,
+                                          value=20,
+                                          label='Number of Steps')
+                    seed = gr.Slider(0,
+                                     100000,
+                                     step=1,
+                                     value=1234,
+                                     label='Seed')
+                    run_button = gr.Button('Run')
+            with gr.Column():
+                result = gr.Image(show_label=False, elem_id='result')
 
         model_name.change(fn=model.set_pipeline,
                           inputs=[
@@ -130,9 +107,6 @@ 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,
@@ -141,10 +115,44 @@ def main():
                              seed,
                          ],
                          outputs=result)
-        model_name2.change(fn=get_sample_image_markdown,
-                           inputs=model_name2,
-                           outputs=sample_images)
+    return demo
 
+
+def create_sample_image_view_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        with gr.Row():
+            model_name = gr.Dropdown([
+                'ddpm-128-exp000 (DDPM)',
+                'ddpm-128-exp000 (DDIM, 20 steps)',
+            ],
+                                     value='ddpm-128-exp000 (DDPM)',
+                                     label='Model')
+        with gr.Row():
+            text = get_sample_image_markdown(model_name.value)
+            sample_images = gr.Markdown(text)
+
+        model_name.change(fn=get_sample_image_markdown,
+                          inputs=model_name,
+                          outputs=sample_images)
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    args = parser.parse_args()
+    model = Model(args.device)
+
+    with gr.Blocks(css='style.css') as demo:
+        gr.Markdown(TITLE)
+        with gr.Tabs():
+            with gr.TabItem('Simple Mode'):
+                create_simple_demo(model)
+            with gr.TabItem('Advanced Mode'):
+                create_advanced_demo(model)
+            with gr.TabItem('Sample Images'):
+                create_sample_image_view_demo()
+        gr.Markdown(FOOTER)
     demo.launch(enable_queue=True, share=False)