download checkpoints first

app.py

@@ -5,47 +5,8 @@ import gradio as gr
 import spaces
 
 PWD = os.path.dirname(__file__)
-# CHECKPOINTS_PATH = "/data/checkpoints"
-CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
-
-# import subprocess
-
-# copy cudnn files
-os.system("cp /root/.pyenv/versions/3.10.18/lib/python3.10/site-packages/nvidia/cudnn/include/*.h /usr/local/cuda/include")
-os.system("cp /root/.pyenv/versions/3.10.18/lib/python3.10/site-packages/nvidia/cudnn/lib/*.so* /usr/local/cuda/lib64")
-
-os.system("apt-get update && apt-get install -qqy libmagickwand-dev")
-
-# install packages
-# os.system('export FLASH_ATTENTION_SKIP_CUDA_BUILD=FALSE && pip install --timeout=1000000000 --no-build-isolation "flash-attn<=2.7.4.post1"')
-os.system(
-    "pip install --timeout=1000000000 https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl"
-)
-os.system('export VLLM_ATTENTION_BACKEND=FLASHINFER && pip install "vllm==0.9.0"')
-os.system('pip install "decord==0.6.0"')
-
-os.system(
-    "export CONDA_PREFIX=/usr/local/cuda && ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/"
-)
-os.system(
-    "export CONDA_PREFIX=/usr/local/cuda && ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10"
-)
-
-os.system('pip install --timeout=1000000000 --no-build-isolation "transformer-engine[pytorch]"')
-os.system('pip install --timeout=1000000000 "decord==0.6.0"')
-
-os.system(
-    'pip install --timeout=1000000000 "git+https://github.com/nvidia-cosmos/cosmos-transfer1@e4055e39ee9c53165e85275bdab84ed20909714a"'
-)
-
-# setup env
-os.environ["CUDA_HOME"] = "/usr/local/cuda"
-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"
-
-from test_environment import main as check_environment
-
-check_environment()
+CHECKPOINTS_PATH = "/data/checkpoints"
+# CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
 
 try:
     import os
@@ -68,12 +29,26 @@ 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 test_environment import main as check_environment
+from test_environment import setup_environment
+
+setup_environment()
+
+# setup env
+os.environ["CUDA_HOME"] = "/usr/local/cuda"
+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"
+
+check_environment()
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Workaround to suppress MP warning
 
 import copy
 import json
 import random
 from io import BytesIO
+
 import torch
 from cosmos_transfer1.checkpoints import (
     BASE_7B_CHECKPOINT_AV_SAMPLE_PATH,

checkpoints.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Cosmos Transfer1 Model Checkpoints
+COSMOS_TRANSFER1_7B_CHECKPOINT = "nvidia/Cosmos-Transfer1-7B"
+COSMOS_TRANSFER1_7B_SAMPLE_AV_CHECKPOINT = "nvidia/Cosmos-Transfer1-7B-Sample-AV"
+COSMOS_TOKENIZER_CHECKPOINT = "nvidia/Cosmos-Tokenize1-CV8x8x8-720p"
+COSMOS_UPSAMPLER_CHECKPOINT = "nvidia/Cosmos-UpsamplePrompt1-12B-Transfer"
+COSMOS_GUARDRAIL_CHECKPOINT = "nvidia/Cosmos-Guardrail1"
+
+# 3rd Party Model Checkpoints
+SAM2_MODEL_CHECKPOINT = "facebook/sam2-hiera-large"
+DEPTH_ANYTHING_MODEL_CHECKPOINT = "depth-anything/Depth-Anything-V2-Small-hf"
+GROUNDING_DINO_MODEL_CHECKPOINT = "IDEA-Research/grounding-dino-tiny"
+T5_MODEL_CHECKPOINT = "google-t5/t5-11b"
+LLAMA_GUARD_3_MODEL_CHECKPOINT = "meta-llama/Llama-Guard-3-8B"
+
+# Internal Checkpoint Paths, please append _PATH to the end of the variable
+BASE_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/base_model.pt"
+VIS2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/vis_control.pt"
+EDGE2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/edge_control.pt"
+SEG2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/seg_control.pt"
+DEPTH2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/depth_control.pt"
+KEYPOINT2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/keypoint_control.pt"
+UPSCALER_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/4kupscaler_control.pt"
+BASE_7B_CHECKPOINT_AV_SAMPLE_PATH = f"{COSMOS_TRANSFER1_7B_SAMPLE_AV_CHECKPOINT}/base_model.pt"
+HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_SAMPLE_AV_CHECKPOINT}/hdmap_control.pt"
+LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_SAMPLE_AV_CHECKPOINT}/lidar_control.pt"
+
+# Transfer1-7B-SV2MV-Sample-AV checkpoints
+
+COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT = "nvidia/Cosmos-Transfer1-7B-Sample-AV-Single2MultiView"
+BASE_t2w_7B_SV2MV_CHECKPOINT_AV_SAMPLE_PATH = f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/t2w_base_model.pt"
+BASE_v2w_7B_SV2MV_CHECKPOINT_AV_SAMPLE_PATH = f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/v2w_base_model.pt"
+SV2MV_t2w_HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/t2w_hdmap_control.pt"
+)
+SV2MV_t2w_LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/t2w_lidar_control.pt"
+)
+SV2MV_v2w_HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/v2w_hdmap_control.pt"
+)
+SV2MV_v2w_LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/v2w_lidar_control.pt"
+)
+
+SV2MV_t2w_HDMAP2WORLD_CONTROLNET_7B_WAYMO_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/t2w_hdmap_control_waymo.pt"
+)
+SV2MV_v2w_HDMAP2WORLD_CONTROLNET_7B_WAYMO_CHECKPOINT_PATH = (
+    f"{COSMOS_TRANSFER1_7B_MV_SAMPLE_AV_CHECKPOINT}/v2w_hdmap_control_waymo.pt"
+)
+EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH = f"{COSMOS_TRANSFER1_7B_CHECKPOINT}/edge_control_distilled.pt"

download_checkpoints.py

@@ -1,11 +1,11 @@
 import hashlib
+import logging
 import os
 import pathlib
 from typing import Literal
 
 # Import the checkpoint paths
-from cosmos_transfer1 import checkpoints
-from cosmos_transfer1.utils import log
+import checkpoints
 from huggingface_hub import login, snapshot_download
 
 
@@ -16,7 +16,7 @@ def download_checkpoint(checkpoint: str, output_dir: str) -> None:
         checkpoint, revision = checkpoint.split(":") if ":" in checkpoint else (checkpoint, None)
         checkpoint_dir = os.path.join(output_dir, checkpoint)
         if get_md5_checksum(output_dir, checkpoint):
-            log.warning(f"Checkpoint {checkpoint_dir} EXISTS, skipping download... ")
+            logging.warning(f"Checkpoint {checkpoint_dir} EXISTS, skipping download... ")
             return
         else:
             print(f"Downloading {checkpoint} to {checkpoint_dir}")
@@ -113,20 +113,20 @@ def main(hf_token: str = os.environ.get("HF_TOKEN"), output_dir: str = "./checkp
             else:
                 checkpoint_vars.append(obj)
 
+    # checkpoint_vars = [
+    #     "nvidia/Cosmos-Guardrail1",
+    #     "nvidia/Cosmos-Tokenize1-CV8x8x8-720p",
+    #     "nvidia/Cosmos-Transfer1-7B-Sample-AV-Single2MultiView",
+    #     "nvidia/Cosmos-Transfer1-7B-Sample-AV",
+    #     "nvidia/Cosmos-UpsamplePrompt1-12B-Transfer",
+    #     "depth-anything/Depth-Anything-V2-Small-hf",
+    #     "IDEA-Research/grounding-dino-tiny",
+    #     "meta-llama/Llama-Guard-3-8B",
+    #     "facebook/sam2-hiera-large",
+    #     "google-t5/t5-11b",
+    # ]
     print(f"Found {len(checkpoint_vars)} checkpoints to download")
     print(checkpoint_vars)
-    checkpoint_vars = [
-        # "nvidia/Cosmos-Guardrail1",
-        # "nvidia/Cosmos-Tokenize1-CV8x8x8-720p",
-        # "nvidia/Cosmos-Transfer1-7B-Sample-AV-Single2MultiView",
-        # "nvidia/Cosmos-Transfer1-7B-Sample-AV",
-        # "nvidia/Cosmos-UpsamplePrompt1-12B-Transfer",
-        # "depth-anything/Depth-Anything-V2-Small-hf",
-        # "IDEA-Research/grounding-dino-tiny",
-        # "meta-llama/Llama-Guard-3-8B",
-        # "facebook/sam2-hiera-large",
-        # "google-t5/t5-11b",
-    ]
 
     # Download each checkpoint
     for checkpoint in checkpoint_vars:

test_environment.py

@@ -19,6 +19,32 @@ import os
 import sys
 
 
+def setup_environment():
+    os.system("apt-get update && apt-get install -qqy libmagickwand-dev")
+
+    # install packages
+    # os.system('export FLASH_ATTENTION_SKIP_CUDA_BUILD=FALSE && pip install --timeout=1000000000 --no-build-isolation "flash-attn<=2.7.4.post1"')
+    os.system(
+        "pip install --timeout=1000000000 https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl"
+    )
+    os.system('export VLLM_ATTENTION_BACKEND=FLASHINFER && pip install "vllm==0.9.0"')
+    os.system('pip install "decord==0.6.0"')
+
+    os.system(
+        "export CONDA_PREFIX=/usr/local/cuda && ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/"
+    )
+    os.system(
+        "export CONDA_PREFIX=/usr/local/cuda && ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10"
+    )
+
+    os.system('pip install --timeout=1000000000 --no-build-isolation "transformer-engine[pytorch]"')
+    os.system('pip install --timeout=1000000000 "decord==0.6.0"')
+
+    os.system(
+        'pip install --timeout=1000000000 "git+https://github.com/nvidia-cosmos/cosmos-transfer1@e4055e39ee9c53165e85275bdab84ed20909714a"'
+    )
+
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(