hysts commited on
Commit
6992dd6
β€’
1 Parent(s): 8aac645

Update for HF Space

Browse files
Files changed (5) hide show
  1. README.md +2 -1
  2. app.py +16 -24
  3. model.py +60 -3
  4. requirements.txt +3 -6
  5. prompt.txt β†’ samples.txt +0 -0
README.md CHANGED
@@ -4,7 +4,8 @@ emoji: 🌍
4
  colorFrom: indigo
5
  colorTo: yellow
6
  sdk: gradio
7
- sdk_version: 3.0.26
 
8
  app_file: app.py
9
  pinned: false
10
  ---
 
4
  colorFrom: indigo
5
  colorTo: yellow
6
  sdk: gradio
7
+ sdk_version: 3.1.0
8
+ python_version: 3.9.13
9
  app_file: app.py
10
  pinned: false
11
  ---
app.py CHANGED
@@ -2,26 +2,20 @@
2
 
3
  from __future__ import annotations
4
 
5
- import argparse
6
-
7
  import gradio as gr
8
 
9
  from model import AppModel
10
 
11
  DESCRIPTION = '''# <a href="https://github.com/THUDM/CogVideo">CogVideo</a>
12
 
13
- The model takes only Chinese as input.
14
- If you check the "Translate to Chinese" checkbox, the app will use the English to Chinese translation results with [this Space](https://huggingface.co/spaces/chinhon/translation_eng2ch) as input.
15
- But the translation model may mistranslate and the results could be poor.
16
- So, it is also a good idea to input the translation results from other translation services.
17
- '''
18
-
19
 
20
- def parse_args() -> argparse.Namespace:
21
- parser = argparse.ArgumentParser()
22
- parser.add_argument('--only-first-stage', action='store_true')
23
- parser.add_argument('--share', action='store_true')
24
- return parser.parse_args()
 
25
 
26
 
27
  def set_example_text(example: list) -> dict:
@@ -29,8 +23,8 @@ def set_example_text(example: list) -> dict:
29
 
30
 
31
  def main():
32
- args = parse_args()
33
- model = AppModel(args.only_first_stage)
34
 
35
  with gr.Blocks(css='style.css') as demo:
36
  gr.Markdown(DESCRIPTION)
@@ -48,14 +42,12 @@ def main():
48
  label='Seed')
49
  only_first_stage = gr.Checkbox(
50
  label='Only First Stage',
51
- value=args.only_first_stage,
52
- visible=not args.only_first_stage)
53
  run_button = gr.Button('Run')
54
 
55
  with open('samples.txt') as f:
56
- samples = [
57
- line.strip().split('\t') for line in f.readlines()
58
- ]
59
  examples = gr.Dataset(components=[text], samples=samples)
60
 
61
  with gr.Column():
@@ -67,6 +59,9 @@ def main():
67
  with gr.TabItem('Output (Gallery)'):
68
  result_gallery = gr.Gallery(show_label=False)
69
 
 
 
 
70
  run_button.click(fn=model.run_with_translation,
71
  inputs=[
72
  text,
@@ -83,10 +78,7 @@ def main():
83
  inputs=examples,
84
  outputs=examples.components)
85
 
86
- demo.launch(
87
- enable_queue=True,
88
- share=args.share,
89
- )
90
 
91
 
92
  if __name__ == '__main__':
 
2
 
3
  from __future__ import annotations
4
 
 
 
5
  import gradio as gr
6
 
7
  from model import AppModel
8
 
9
  DESCRIPTION = '''# <a href="https://github.com/THUDM/CogVideo">CogVideo</a>
10
 
11
+ Currently, this Space only supports the first stage of the CogVideo pipeline due to hardware limitations.
 
 
 
 
 
12
 
13
+ The model accepts only Chinese as input.
14
+ 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.
15
+ Since the translation model may mistranslate, you may want to use the translation results from other translation services.
16
+ '''
17
+ NOTES = 'This app is adapted from https://github.com/hysts/CogVideo_demo. It would be recommended to use the repo if you want to run the app yourself.'
18
+ FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=THUDM.CogVideo" />'
19
 
20
 
21
  def set_example_text(example: list) -> dict:
 
23
 
24
 
25
  def main():
26
+ only_first_stage = True
27
+ model = AppModel(only_first_stage)
28
 
29
  with gr.Blocks(css='style.css') as demo:
30
  gr.Markdown(DESCRIPTION)
 
42
  label='Seed')
43
  only_first_stage = gr.Checkbox(
44
  label='Only First Stage',
45
+ value=only_first_stage,
46
+ visible=not only_first_stage)
47
  run_button = gr.Button('Run')
48
 
49
  with open('samples.txt') as f:
50
+ samples = [[line.strip()] for line in f.readlines()]
 
 
51
  examples = gr.Dataset(components=[text], samples=samples)
52
 
53
  with gr.Column():
 
59
  with gr.TabItem('Output (Gallery)'):
60
  result_gallery = gr.Gallery(show_label=False)
61
 
62
+ gr.Markdown(NOTES)
63
+ gr.Markdown(FOOTER)
64
+
65
  run_button.click(fn=model.run_with_translation,
66
  inputs=[
67
  text,
 
78
  inputs=examples,
79
  outputs=examples.components)
80
 
81
+ demo.launch(enable_queue=True, share=False)
 
 
 
82
 
83
 
84
  if __name__ == '__main__':
model.py CHANGED
@@ -1,16 +1,72 @@
1
- # This code is adapted from https://github.com/THUDM/CogView2/blob/4e55cce981eb94b9c8c1f19ba9f632fd3ee42ba8/cogview2_text2image.py
2
 
3
  from __future__ import annotations
4
 
5
  import argparse
6
- import functools
7
  import logging
 
8
  import pathlib
 
 
9
  import sys
10
  import tempfile
11
  import time
 
12
  from typing import Any
13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  import gradio as gr
15
  import imageio.v2 as iio
16
  import numpy as np
@@ -1116,6 +1172,7 @@ class Model:
1116
  start = time.perf_counter()
1117
 
1118
  set_random_seed(seed)
 
1119
 
1120
  if only_first_stage:
1121
  self.args.stage_1 = True
@@ -1169,7 +1226,7 @@ class AppModel(Model):
1169
 
1170
  def run_with_translation(
1171
  self, text: str, translate: bool, seed: int, only_first_stage: bool
1172
- ) -> tuple[str | None, np.ndarray | None, list[np.ndarray] | None]:
1173
  logger.info(f'{text=}, {translate=}, {seed=}, {only_first_stage=}')
1174
  if translate:
1175
  text = translated_text = self.translator(text)
 
1
+ # This code is adapted from https://github.com/THUDM/CogVideo/blob/ff423aa169978fb2f636f761e348631fa3178b03/cogvideo_pipeline.py
2
 
3
  from __future__ import annotations
4
 
5
  import argparse
 
6
  import logging
7
+ import os
8
  import pathlib
9
+ import shutil
10
+ import subprocess
11
  import sys
12
  import tempfile
13
  import time
14
+ import zipfile
15
  from typing import Any
16
 
17
+ if os.getenv('SYSTEM') == 'spaces':
18
+ subprocess.run('pip install icetk==0.0.4'.split())
19
+ subprocess.run('pip install SwissArmyTransformer==0.2.9'.split())
20
+ subprocess.run(
21
+ 'pip install git+https://github.com/Sleepychord/Image-Local-Attention@43fee31'
22
+ .split())
23
+ #subprocess.run('git clone https://github.com/NVIDIA/apex'.split())
24
+ #subprocess.run('git checkout 1403c21'.split(), cwd='apex')
25
+ #with open('patch.apex') as f:
26
+ # subprocess.run('patch -p1'.split(), cwd='apex', stdin=f)
27
+ #subprocess.run(
28
+ # 'pip install -v --disable-pip-version-check --no-cache-dir --global-option --cpp_ext --global-option --cuda_ext ./'
29
+ # .split(),
30
+ # cwd='apex')
31
+ #subprocess.run('rm -rf apex'.split())
32
+ with open('patch') as f:
33
+ subprocess.run('patch -p1'.split(), cwd='CogVideo', stdin=f)
34
+
35
+ from huggingface_hub import hf_hub_download
36
+
37
+ def download_and_extract_icetk_models() -> None:
38
+ icetk_model_dir = pathlib.Path('/home/user/.icetk_models')
39
+ icetk_model_dir.mkdir()
40
+ path = hf_hub_download('THUDM/icetk',
41
+ 'models.zip',
42
+ use_auth_token=os.getenv('HF_TOKEN'))
43
+ with zipfile.ZipFile(path) as f:
44
+ f.extractall(path=icetk_model_dir.as_posix())
45
+
46
+ def download_and_extract_cogvideo_models(name: str) -> None:
47
+ path = hf_hub_download('THUDM/CogVideo',
48
+ name,
49
+ use_auth_token=os.getenv('HF_TOKEN'))
50
+ with zipfile.ZipFile(path) as f:
51
+ f.extractall('pretrained')
52
+ os.remove(path)
53
+
54
+ def download_and_extract_cogview2_models(name: str) -> None:
55
+ path = hf_hub_download('THUDM/CogView2', name)
56
+ with zipfile.ZipFile(path) as f:
57
+ f.extractall()
58
+ shutil.move('/home/user/app/sharefs/cogview-new/cogview2-dsr',
59
+ 'pretrained')
60
+ shutil.rmtree('/home/user/app/sharefs/')
61
+ os.remove(path)
62
+
63
+ download_and_extract_icetk_models()
64
+ download_and_extract_cogvideo_models('cogvideo-stage1.zip')
65
+ #download_and_extract_cogvideo_models('cogvideo-stage2.zip')
66
+ #download_and_extract_cogview2_models('cogview2-dsr.zip')
67
+
68
+ os.environ['SAT_HOME'] = '/home/user/app/pretrained'
69
+
70
  import gradio as gr
71
  import imageio.v2 as iio
72
  import numpy as np
 
1172
  start = time.perf_counter()
1173
 
1174
  set_random_seed(seed)
1175
+ self.args.seed = seed
1176
 
1177
  if only_first_stage:
1178
  self.args.stage_1 = True
 
1226
 
1227
  def run_with_translation(
1228
  self, text: str, translate: bool, seed: int, only_first_stage: bool
1229
+ ) -> tuple[str | None, str | None, list[np.ndarray] | None]:
1230
  logger.info(f'{text=}, {translate=}, {seed=}, {only_first_stage=}')
1231
  if translate:
1232
  text = translated_text = self.translator(text)
requirements.txt CHANGED
@@ -1,10 +1,7 @@
1
- git+https://github.com/Sleepychord/Image-Local-Attention@43fee31
2
- gradio==3.1.0
3
- icetk==0.0.4
4
  imageio==2.19.5
5
  imageio-ffmpeg==0.4.7
6
  numpy==1.22.4
7
  opencv-python-headless==4.6.0.66
8
- SwissArmyTransformer==0.2.9
9
- torch==1.12.0
10
- torchvision==0.13.0
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
 
 
2
  imageio==2.19.5
3
  imageio-ffmpeg==0.4.7
4
  numpy==1.22.4
5
  opencv-python-headless==4.6.0.66
6
+ torch==1.12.0+cu113
7
+ torchvision==0.13.0+cu113
 
prompt.txt β†’ samples.txt RENAMED
File without changes