move assignment of CHECKPOINTS_PATH to `app.py`

app.py

@@ -31,15 +31,17 @@ if not check_environment():
 
 
 from download_checkpoints import main as download_checkpoints
-from helper import CHECKPOINTS_PATH, generate_video
+
+CHECKPOINTS_PATH = "/data/checkpoints"
 
 # download checkpoints
 os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
 download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
 
-
 import gradio as gr
 
+from helper import generate_video_fun
+
 # Define the Gradio Blocks interface
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -81,7 +83,7 @@ with gr.Blocks() as demo:
             output_file = gr.File(label="Download Results")
 
     generate_button.click(
-        fn=generate_video,
+        fn=generate_video_fun(CHECKPOINTS_PATH),
         inputs=[
             rgb_video_input,
             hdmap_input,

download_checkpoints.py

@@ -131,3 +131,12 @@ def main(hf_token: str = os.environ.get("HF_TOKEN"), output_dir: str = "./checkp
     # Download each checkpoint
     for checkpoint in checkpoint_vars:
         download_checkpoint(checkpoint, output_dir)
+
+
+if __name__ == "__main__":
+    import sys
+
+    PWD = os.path.dirname(__file__)
+    CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
+    os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+    main(hf_token=sys.argv[1], output_dir=CHECKPOINTS_PATH, model="all")

helper.py

@@ -32,7 +32,6 @@ from cosmos_transfer1.utils.io import read_prompts_from_file, save_video
 from gpu_info import stop_watcher, watch_gpu_memory
 
 PWD = os.path.dirname(__file__)
-CHECKPOINTS_PATH = "/data/checkpoints"
 LOG_DIR = os.path.join(PWD, "logs")
 os.makedirs(LOG_DIR, exist_ok=True)
 
@@ -370,68 +369,71 @@ def create_zip_for_download(filename, files_to_zip):
 import gradio as gr
 
 
-def generate_video(
-    rgb_video_path,
-    hdmap_video_input,
-    lidar_video_input,
-    prompt,
-    negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.",  # noqa: E501
-    seed=42,
-    randomize_seed=False,
-    chunking=None,
-    progress=gr.Progress(track_tqdm=True),
-):
-    _dt = datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=8))).strftime("%Y-%m-%d_%H.%M.%S")
-    logfile_path = os.path.join(LOG_DIR, f"{_dt}.log")
-    log_handler = log.init_dev_loguru_file(logfile_path)
-
-    if randomize_seed:
-        actual_seed = random.randint(0, 1000000)
-    else:
-        actual_seed = seed
-
-    log.info(f"actual_seed: {actual_seed}")
-
-    if rgb_video_path is None or not os.path.isfile(rgb_video_path):
-        log.warning(f"File `{rgb_video_path}` does not exist")
-        rgb_video_path = ""
-
-    # add timer to calculate the generation time
-    start_time = time.time()
+def generate_video_fun(checkpoints_path: str):
+    def generate_video(
+        rgb_video_path,
+        hdmap_video_input,
+        lidar_video_input,
+        prompt,
+        negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.",  # noqa: E501
+        seed=42,
+        randomize_seed=False,
+        chunking=None,
+        progress=gr.Progress(track_tqdm=True),
+    ):
+        _dt = datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=8))).strftime("%Y-%m-%d_%H.%M.%S")
+        logfile_path = os.path.join(LOG_DIR, f"{_dt}.log")
+        log_handler = log.init_dev_loguru_file(logfile_path)
+
+        if randomize_seed:
+            actual_seed = random.randint(0, 1000000)
+        else:
+            actual_seed = seed
+
+        log.info(f"actual_seed: {actual_seed}")
+
+        if rgb_video_path is None or not os.path.isfile(rgb_video_path):
+            log.warning(f"File `{rgb_video_path}` does not exist")
+            rgb_video_path = ""
+
+        # add timer to calculate the generation time
+        start_time = time.time()
+
+        # parse generation configs
+        args, control_inputs = parse_arguments(
+            controlnet_specs_in={
+                "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
+                "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
+            },
+            input_video_path=rgb_video_path,
+            checkpoint_dir=checkpoints_path,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            sigma_max=80,
+            offload_text_encoder_model=True,
+            is_av_sample=True,
+            num_gpus=1,
+            seed=seed,
+        )
 
-    # parse generation configs
-    args, control_inputs = parse_arguments(
-        controlnet_specs_in={
-            "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
-            "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
-        },
-        input_video_path=rgb_video_path,
-        checkpoint_dir=CHECKPOINTS_PATH,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        sigma_max=80,
-        offload_text_encoder_model=True,
-        is_av_sample=True,
-        num_gpus=1,
-        seed=seed,
-    )
+        # watch gpu memory
+        watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory (used, total): {x} (MiB)"))
 
-    # watch gpu memory
-    watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory (used, total): {x} (MiB)"))
+        # start inference
+        if chunking <= 0:
+            chunking = None
+        videos, prompts = inference(args, control_inputs, chunking)
 
-    # start inference
-    if chunking <= 0:
-        chunking = None
-    videos, prompts = inference(args, control_inputs, chunking)
+        # print the generation time
+        end_time = time.time()
+        log.info(f"Time taken: {end_time - start_time} s")
 
-    # print the generation time
-    end_time = time.time()
-    log.info(f"Time taken: {end_time - start_time} s")
+        # stop the watcher
+        stop_watcher()
 
-    # stop the watcher
-    stop_watcher()
+        video = videos[0]
 
-    video = videos[0]
+        log.logger.remove(log_handler)
+        return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
 
-    log.logger.remove(log_handler)
-    return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
+    return generate_video

main.py

@@ -1,11 +1,11 @@
-from argparse import ArgumentParser
 import os
 import sys
-
+from argparse import ArgumentParser
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    parser.add_argument("--skip-check-environment", action="store_true", help="Whether to skip environment checking")
+    parser.add_argument("--check-environment", action="store_true", help="Whether to do the environment checking")
+    parser.add_argument("--download-checkpoints", action="store_true", help="Whether to download the checkpoints")
     parser.add_argument("--rgb", type=str, default=None, help="RGB video input")
     parser.add_argument("--hdmap", type=str, required=True, help="HD Map video input")
     parser.add_argument("--lidar", type=str, required=True, help="Lidar video input")
@@ -31,22 +31,28 @@ if __name__ == "__main__":
     os.environ["LD_LIBRARY_PATH"] = "$CUDA_HOME/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH"
     os.environ["PATH"] = "$CUDA_HOME/bin:/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:$PATH"
 
-    if not args.skip_check_environment:
+    if args.check_environment:
         from test_environment import main as check_environment
 
         if not check_environment():
             sys.exit(1)
     else:
-        print("Skipping environment check")
+        print("Skipping the environment checking")
 
-    from download_checkpoints import main as download_checkpoints
-    from helper import CHECKPOINTS_PATH, generate_video
+    from helper import PWD
+
+    CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
 
     # download checkpoints
-    os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
-    download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
+    if args.download_checkpoints:
+        from download_checkpoints import main as download_checkpoints
+
+        os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+        download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
+
+    from helper import generate_video_fun
 
-    generate_video(
+    generate_video_fun(CHECKPOINTS_PATH)(
         rgb_video_path=args.rgb,
         hdmap_video_input=args.hdmap,
         lidar_video_input=args.lidar,