File size: 3,195 Bytes
649bf38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
#!/usr/bin/env python
from __future__ import annotations
import argparse
import gradio as gr
from model import AppModel
DESCRIPTION = '''# <a href="https://github.com/THUDM/CogVideo">CogVideo</a>
The model accepts only Chinese as 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.
'''
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument('--only-first-stage', action='store_true')
parser.add_argument('--share', action='store_true')
return parser.parse_args()
def set_example_text(example: list) -> dict:
return gr.Textbox.update(value=example[0])
def main():
args = parse_args()
model = AppModel(args.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)
seed = gr.Slider(0,
100000,
step=1,
value=1234,
label='Seed')
only_first_stage = gr.Checkbox(
label='Only First Stage',
value=args.only_first_stage,
visible=not args.only_first_stage)
run_button = gr.Button('Run')
with open('samples.txt') as f:
samples = [[line.strip()] for line in f.readlines()]
examples = gr.Dataset(components=[text], samples=samples)
with gr.Column():
with gr.Group():
translated_text = gr.Textbox(label='Translated Text')
with gr.Tabs():
with gr.TabItem('Output (Video)'):
result_video = gr.Video(show_label=False)
with gr.TabItem('Output (Gallery)'):
result_gallery = gr.Gallery(show_label=False)
run_button.click(fn=model.run_with_translation,
inputs=[
text,
translate,
seed,
only_first_stage,
],
outputs=[
translated_text,
result_video,
result_gallery,
])
examples.click(fn=set_example_text,
inputs=examples,
outputs=examples.components)
demo.launch(
enable_queue=True,
share=args.share,
)
if __name__ == '__main__':
main()
|