change checkpoints path

app.py

@@ -5,26 +5,33 @@ import gradio as gr
 import spaces
 
 PWD = os.path.dirname(__file__)
+CHECKPOINTS_PATH = "/data/checkpoints"
 
-import subprocess
+# import subprocess
 
 # copy cudnn files
-subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/include/*.h /usr/local/cuda/include", env={}, shell=True)
-subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/*.so* /usr/local/cuda/lib64", env={}, shell=True)
+# subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/include/*.h /usr/local/cuda/include", env={}, shell=True)
+# subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/*.so* /usr/local/cuda/lib64", env={}, shell=True)
 
 # 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"
+# 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"
+
+os.system("apt-get update && apt-get install -qqy libmagickwand-dev")
 
 # install packages
-subprocess.run("pip install flash-attn --no-build-isolation", env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}, shell=True)
-subprocess.run("pip install vllm==0.9.0", env={"VLLM_ATTENTION_BACKEND": "FLASHINFER"}, shell=True)
-subprocess.run(
-    "pip install transformer-engine[pytorch] --no-build-isolation",
-    env={"PATH": os.environ["PATH"], "LD_LIBRARY_PATH": os.environ["LD_LIBRARY_PATH"], "CUDA_HOME": os.environ["CUDA_HOME"]},
-    shell=True,
+os.system('export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE && pip install "flash-attn<=2.7.4.post1" --no-build-isolation')
+os.system(
+    "pip install 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("ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/")
+os.system("ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10")
+
+os.system('pip install "transformer-engine[pytorch]"')
 
 from test_environment import main as check_environment
 
@@ -48,8 +55,8 @@ except Exception as e:
 # download checkpoints
 from download_checkpoints import main as download_checkpoints
 
-os.makedirs(os.path.join(PWD, "checkpoints"), exist_ok=True)
-download_checkpoints(hf_token="", output_dir=os.path.join(PWD, "checkpoints"), model="7b_av")
+os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Workaround to suppress MP warning
 
@@ -292,7 +299,7 @@ def generate_video(
             "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
             "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
         },
-        checkpoint_dir=os.path.join(PWD, "checkpoints"),
+        checkpoint_dir=CHECKPOINTS_PATH,
         prompt=prompt,
         negative_prompt=negative_prompt,
         sigma_max=80,

download_checkpoints.py

@@ -115,6 +115,18 @@ def main(hf_token: str = os.environ.get("HF_TOKEN"), output_dir: str = "./checkp
 
     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:

requirements.txt

@@ -1,12 +1,7 @@
 # essentials
+wheel
 ninja
 
-# torch
-# torch==2.7.0
-# torchvision==0.22.0
-# torch==2.5.1
-# torchvision==0.20.1
-
 # cosmos-transfer1
 git+https://github.com/nvidia-cosmos/cosmos-transfer1