Add a simpler interface

app.py

@@ -6,8 +6,9 @@ import gradio as gr
 
 from model import AppModel
 
-DESCRIPTION = '''# <a href="https://github.com/THUDM/CogView2">CogView2</a> (text2image)
+TITLE = '# <a href="https://github.com/THUDM/CogView2">CogView2</a> (text2image)'
 
+DESCRIPTION = '''
 The model accepts English or Chinese as input.
 In general, Chinese input produces better results than English input.
 By checking the "Translate to Chinese" checkbox, the results of English to Chinese translation with [this Space](https://huggingface.co/spaces/chinhon/translation_eng2ch) will be used as input. Since the translation model may mistranslate, you may want to use the translation results from other translation services.
@@ -31,64 +32,94 @@ def main():
     model = AppModel(max_inference_batch_size, only_first_stage)
 
     with gr.Blocks(css='style.css') as demo:
-        gr.Markdown(DESCRIPTION)
-
-        with gr.Row():
-            with gr.Column():
-                with gr.Group():
-                    text = gr.Textbox(label='Input Text')
-                    translate = gr.Checkbox(label='Translate to Chinese',
-                                            value=False)
-                    style = gr.Dropdown(choices=[
-                        'none',
-                        'mainbody',
-                        'photo',
-                        'flat',
-                        'comics',
-                        'oil',
-                        'sketch',
-                        'isometric',
-                        'chinese',
-                        'watercolor',
-                    ],
-                                        value='mainbody',
-                                        label='Style')
-                    seed = gr.Slider(0,
-                                     100000,
-                                     step=1,
-                                     value=1234,
-                                     label='Seed')
-                    only_first_stage = gr.Checkbox(
-                        label='Only First Stage',
-                        value=only_first_stage,
-                        visible=not only_first_stage)
-                    num_images = gr.Slider(1,
-                                           16,
-                                           step=1,
-                                           value=4,
-                                           label='Number of Images')
-                    run_button = gr.Button('Run')
-
-                    with open('samples.txt') as f:
-                        samples = [
-                            line.strip().split('\t') for line in f.readlines()
-                        ]
-                    examples = gr.Dataset(components=[text, style],
-                                          samples=samples)
-
-            with gr.Column():
-                with gr.Group():
-                    translated_text = gr.Textbox(label='Translated Text')
-                    with gr.Tabs():
-                        with gr.TabItem('Output (Grid View)'):
-                            result_grid = gr.Image(show_label=False)
-                        with gr.TabItem('Output (Gallery)'):
-                            result_gallery = gr.Gallery(show_label=False)
-
-        gr.Markdown(NOTES)
+
+        with gr.Tabs():
+            with gr.TabItem('Simple Mode'):
+                gr.Markdown(TITLE)
+
+                with gr.Row().style(mobile_collapse=False, equal_height=True):
+                    text_simple = gr.Textbox(placeholder='Enter your prompt',
+                                             show_label=False,
+                                             max_lines=1).style(
+                                                 border=(True, False, True,
+                                                         True),
+                                                 rounded=(True, False, False,
+                                                          True),
+                                                 container=False,
+                                             )
+                    run_button_simple = gr.Button('Run').style(
+                        margin=False,
+                        rounded=(False, True, True, False),
+                    )
+                result_grid_simple = gr.Image(show_label=False)
+
+            with gr.TabItem('Advanced Mode'):
+                gr.Markdown(TITLE)
+                gr.Markdown(DESCRIPTION)
+
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Group():
+                            text = gr.Textbox(label='Input Text')
+                            translate = gr.Checkbox(
+                                label='Translate to Chinese', value=False)
+                            style = gr.Dropdown(choices=[
+                                'none',
+                                'mainbody',
+                                'photo',
+                                'flat',
+                                'comics',
+                                'oil',
+                                'sketch',
+                                'isometric',
+                                'chinese',
+                                'watercolor',
+                            ],
+                                                value='mainbody',
+                                                label='Style')
+                            seed = gr.Slider(0,
+                                             100000,
+                                             step=1,
+                                             value=1234,
+                                             label='Seed')
+                            only_first_stage = gr.Checkbox(
+                                label='Only First Stage',
+                                value=only_first_stage,
+                                visible=not only_first_stage)
+                            num_images = gr.Slider(1,
+                                                   16,
+                                                   step=1,
+                                                   value=4,
+                                                   label='Number of Images')
+                            run_button = gr.Button('Run')
+
+                            with open('samples.txt') as f:
+                                samples = [
+                                    line.strip().split('\t')
+                                    for line in f.readlines()
+                                ]
+                            examples = gr.Dataset(components=[text, style],
+                                                  samples=samples)
+
+                    with gr.Column():
+                        with gr.Group():
+                            translated_text = gr.Textbox(
+                                label='Translated Text')
+                            with gr.Tabs():
+                                with gr.TabItem('Output (Grid View)'):
+                                    result_grid = gr.Image(show_label=False)
+                                with gr.TabItem('Output (Gallery)'):
+                                    result_gallery = gr.Gallery(
+                                        show_label=False)
+
+                gr.Markdown(NOTES)
+
         gr.Markdown(FOOTER)
 
-        run_button.click(fn=model.run_with_translation,
+        run_button_simple.click(fn=model.run_simple,
+                                inputs=text_simple,
+                                outputs=result_grid_simple)
+        run_button.click(fn=model.run_advanced,
                          inputs=[
                              text,
                              translate,
@@ -104,7 +135,8 @@ def main():
                          ])
         examples.click(fn=set_example_text,
                        inputs=examples,
-                       outputs=examples.components)
+                       outputs=examples.components,
+                       queue=False)
 
     demo.launch(enable_queue=True)
 

model.py

@@ -7,6 +7,7 @@ import functools
 import logging
 import os
 import pathlib
+import random
 import subprocess
 import sys
 import time
@@ -402,6 +403,7 @@ class AppModel(Model):
         super().__init__(max_inference_batch_size, only_first_stage)
         self.translator = gr.Interface.load(
             'spaces/chinhon/translation_eng2ch')
+        self.rng = random.Random()
 
     def make_grid(self, images: list[np.ndarray] | None) -> np.ndarray | None:
         if images is None or len(images) == 0:
@@ -422,7 +424,7 @@ class AppModel(Model):
                 grid[h * i:h * (i + 1), w * j:w * (j + 1)] = images[index]
         return grid
 
-    def run_with_translation(
+    def run_advanced(
         self, text: str, translate: bool, style: str, seed: int,
         only_first_stage: bool, num: int
     ) -> tuple[str | None, np.ndarray | None, list[np.ndarray] | None]:
@@ -436,3 +438,12 @@ class AppModel(Model):
         results = self.run(text, style, seed, only_first_stage, num)
         grid_image = self.make_grid(results)
         return translated_text, grid_image, results
+
+    def run_simple(self, text: str) -> np.ndarray | None:
+        logger.info(f'{text=}')
+        if text.isascii():
+            text = self.translator(text)
+        seed = self.rng.randint(0, 100000)
+        results = self.run(text, 'photo', seed, False, 4)
+        grid_image = self.make_grid(results)
+        return grid_image