Add stage2 support

app.py

@@ -8,7 +8,8 @@ from model import AppModel
 
 DESCRIPTION = '''# <a href="https://github.com/THUDM/CogVideo">CogVideo</a>
 
-Currently, this Space only supports the first stage of the CogVideo pipeline due to hardware limitations.
+This Space supports the first stage and the second stage (better quality) of the CogVideo pipeline. 
+Note that the second stage of CogVideo is **much slower**. 
 
 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.
@@ -19,7 +20,7 @@ FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge
 
 
 def main():
-    only_first_stage = True
+    only_first_stage = False
     model = AppModel(only_first_stage)
 
     with gr.Blocks(css='style.css') as demo:
@@ -36,9 +37,9 @@ def main():
                                      step=1,
                                      value=1234,
                                      label='Seed')
-                    only_first_stage = gr.Checkbox(
+                    run_only_first_stage = gr.Checkbox(
                         label='Only First Stage',
-                        value=only_first_stage,
+                        value=True,
                         visible=not only_first_stage)
                     image_prompt = gr.Image(type="filepath",
                                             label="Image Prompt",
@@ -53,10 +54,10 @@ def main():
                             result_video = gr.Video(show_label=False)
 
         examples = gr.Examples(
-            examples=[['骑滑板的皮卡丘', False, 1234, True,None],
-                      ['a cat playing chess', True, 1253, True,None]],
+            examples=[['骑滑板的皮卡丘', False, 1234, True, None],
+                      ['a cat playing chess', True, 1253, False, None]],
             fn=model.run_with_translation,
-            inputs=[text, translate, seed, only_first_stage,image_prompt],
+            inputs=[text, translate, seed, run_only_first_stage, image_prompt],
             outputs=[translated_text, result_video],
             cache_examples=True)
 
@@ -68,7 +69,7 @@ def main():
                              text,
                              translate,
                              seed,
-                             only_first_stage,
+                             run_only_first_stage,
                              image_prompt
                          ],
                          outputs=[translated_text, result_video])

model.py

@@ -62,8 +62,8 @@ if os.getenv('SYSTEM') == 'spaces':
 
     download_and_extract_icetk_models()
     download_and_extract_cogvideo_models('cogvideo-stage1.zip')
-    #download_and_extract_cogvideo_models('cogvideo-stage2.zip')
-    #download_and_extract_cogview2_models('cogview2-dsr.zip')
+    download_and_extract_cogvideo_models('cogvideo-stage2.zip')
+    download_and_extract_cogview2_models('cogview2-dsr.zip')
 
     os.environ['SAT_HOME'] = '/home/user/app/pretrained'
 
@@ -299,7 +299,8 @@ def my_filling_sequence(
     # initialize generation
     counter = context_length - 1  # Last fixed index is ``counter''
     index = 0  # Next forward starting index, also the length of cache.
-    mems_buffers_on_GPU = False
+    # mems_buffers_on_GPU = False
+    torch.cuda.empty_cache()
     mems_indexs = [0, 0]
     mems_len = [(400 + 74) if limited_spatial_channel_mem else 5 * 400 + 74,
                 5 * 400 + 74]
@@ -308,7 +309,8 @@ def my_filling_sequence(
                     batch_size,
                     mem_len,
                     args.hidden_size * 2,
-                    dtype=next(model.parameters()).dtype)
+                    dtype=next(model.parameters()).dtype,
+                    device=next(model.parameters()).device)
         for mem_len in mems_len
     ]
 
@@ -320,13 +322,13 @@ def my_filling_sequence(
                         batch_size,
                         mem_len,
                         args.hidden_size * 2,
-                        dtype=next(model.parameters()).dtype)
+                        dtype=next(model.parameters()).dtype,
+                        device=next(model.parameters()).device)
             for mem_len in mems_len
         ]
         guider_mems_indexs = [0, 0]
         guider_mems = None
 
-    torch.cuda.empty_cache()
     # step-by-step generation
     while counter < len(seq[0]) - 1:
         # we have generated counter+1 tokens
@@ -448,34 +450,34 @@ def my_filling_sequence(
                 ]
                 guider_logits = guider_logits_all
         else:
-            if not mems_buffers_on_GPU:
-                if not mode_stage1:
-                    torch.cuda.empty_cache()
-                    for idx, mem in enumerate(mems):
-                        mems[idx] = mem.to(next(model.parameters()).device)
-                    if guider_seq is not None:
-                        for idx, mem in enumerate(guider_mems):
-                            guider_mems[idx] = mem.to(
-                                next(model.parameters()).device)
-                else:
-                    torch.cuda.empty_cache()
-                    for idx, mem_buffer in enumerate(mems_buffers):
-                        mems_buffers[idx] = mem_buffer.to(
-                            next(model.parameters()).device)
-                    mems = [
-                        mems_buffers[id][:, :, :mems_indexs[id]]
-                        for id in range(2)
-                    ]
-                    if guider_seq is not None:
-                        for idx, guider_mem_buffer in enumerate(
-                                guider_mems_buffers):
-                            guider_mems_buffers[idx] = guider_mem_buffer.to(
-                                next(model.parameters()).device)
-                        guider_mems = [
-                            guider_mems_buffers[id]
-                            [:, :, :guider_mems_indexs[id]] for id in range(2)
-                        ]
-                    mems_buffers_on_GPU = True
+            # if not mems_buffers_on_GPU:
+            #     if not mode_stage1:
+            #         torch.cuda.empty_cache()
+            #         for idx, mem in enumerate(mems):
+            #             mems[idx] = mem.to(next(model.parameters()).device)
+            #         if guider_seq is not None:
+            #             for idx, mem in enumerate(guider_mems):
+            #                 guider_mems[idx] = mem.to(
+            #                     next(model.parameters()).device)
+            #     else:
+            #         torch.cuda.empty_cache()
+            #         for idx, mem_buffer in enumerate(mems_buffers):
+            #             mems_buffers[idx] = mem_buffer.to(
+            #                 next(model.parameters()).device)
+            #         mems = [
+            #             mems_buffers[id][:, :, :mems_indexs[id]]
+            #             for id in range(2)
+            #         ]
+            #         if guider_seq is not None:
+            #             for idx, guider_mem_buffer in enumerate(
+            #                     guider_mems_buffers):
+            #                 guider_mems_buffers[idx] = guider_mem_buffer.to(
+            #                     next(model.parameters()).device)
+            #             guider_mems = [
+            #                 guider_mems_buffers[id]
+            #                 [:, :, :guider_mems_indexs[id]] for id in range(2)
+            #             ]
+            #         mems_buffers_on_GPU = True
 
             logits, *output_per_layers = model(
                 input_tokens[:, index:],
@@ -513,17 +515,17 @@ def my_filling_sequence(
                     o['mem_kv'][0] for o in guider_output_per_layers
                 ], [o['mem_kv'][1] for o in guider_output_per_layers]
 
-            if not mems_buffers_on_GPU:
-                torch.cuda.empty_cache()
-                for idx, mem_buffer in enumerate(mems_buffers):
-                    mems_buffers[idx] = mem_buffer.to(
-                        next(model.parameters()).device)
-                if guider_seq is not None:
-                    for idx, guider_mem_buffer in enumerate(
-                            guider_mems_buffers):
-                        guider_mems_buffers[idx] = guider_mem_buffer.to(
-                            next(model.parameters()).device)
-                mems_buffers_on_GPU = True
+            # if not mems_buffers_on_GPU:
+            #     torch.cuda.empty_cache()
+            #     for idx, mem_buffer in enumerate(mems_buffers):
+            #         mems_buffers[idx] = mem_buffer.to(
+            #             next(model.parameters()).device)
+            #     if guider_seq is not None:
+            #         for idx, guider_mem_buffer in enumerate(
+            #                 guider_mems_buffers):
+            #             guider_mems_buffers[idx] = guider_mem_buffer.to(
+            #                 next(model.parameters()).device)
+            #     mems_buffers_on_GPU = True
 
             mems, mems_indexs = my_update_mems([mem_kv0, mem_kv1],
                                                mems_buffers, mems_indexs,
@@ -677,7 +679,7 @@ def get_default_args() -> argparse.Namespace:
         '--batch-size',
         '1',
         '--max-inference-batch-size',
-        '8',
+        '1',
     ]
     args = get_args(args_list)
     args = argparse.Namespace(**vars(args), **vars(known))
@@ -779,7 +781,7 @@ class Model:
             path = auto_create('cogview2-dsr', path=None)
             dsr = DirectSuperResolution(self.args,
                                         path,
-                                        max_bz=12,
+                                        max_bz=4,
                                         onCUDA=False)
         else:
             dsr = None
@@ -1184,7 +1186,8 @@ class Model:
         else:
             self.args.stage_1 = False
             self.args.both_stages = True
-
+        
+        torch.cuda.empty_cache()
         parent_given_tokens, res = self.process_stage1(
             self.model_stage1,
             text,
@@ -1231,7 +1234,7 @@ class AppModel(Model):
 
     def run_with_translation(
             self, text: str, translate: bool, seed: int,
-            only_first_stage: bool,image_prompt: None) -> tuple[str | None, str | None]:
+            only_first_stage: bool, image_prompt: None) -> tuple[str | None, str | None]:
             
         logger.info(f'{text=}, {translate=}, {seed=}, {only_first_stage=},{image_prompt=}')
         if translate: