supported downloading ckpt

README.md

@@ -1,5 +1,5 @@
 
-![Cosmos Logo](assets/cosmos-logo.png)
+![Cosmos Logo](https://github.com/NVIDIA/Cosmos/raw/main/assets/cosmos-logo.png)
 
 --------------------------------------------------------------------------------
 ### [Website](https://www.nvidia.com/en-us/ai/cosmos/) | [HuggingFace](https://huggingface.co/collections/nvidia/cosmos-6751e884dc10e013a0a0d8e6) | [GPU-free Preview](https://build.nvidia.com/explore/discover) | [Paper](https://arxiv.org/abs/2501.03575) | [Paper Website](https://research.nvidia.com/labs/dir/cosmos1/)

cosmos1/scripts/download_diffusion.py → download_diffusion.py

@@ -18,7 +18,7 @@ from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
-from cosmos1.scripts.convert_pixtral_ckpt import convert_pixtral_checkpoint
+from .convert_pixtral_ckpt import convert_pixtral_checkpoint
 
 
 def parse_args():
@@ -57,7 +57,7 @@ def parse_args():
     return args
 
 
-def main(args):
+def main(model_types, model_sizes, checkpoint_dir="checkpoints"):
     ORG_NAME = "nvidia"
 
     # Mapping from size argument to Hugging Face repository name
@@ -72,18 +72,18 @@ def main(args):
         "Cosmos-1.0-Tokenizer-CV8x8x8",
     ]
 
-    if "Text2World" in args.model_types:
+    if "Text2World" in model_types:
         extra_models.append("Cosmos-1.0-Prompt-Upsampler-12B-Text2World")
 
     # Create local checkpoints folder
-    checkpoints_dir = Path(args.checkpoint_dir)
+    checkpoints_dir = Path(checkpoint_dir)
     checkpoints_dir.mkdir(parents=True, exist_ok=True)
 
     download_kwargs = dict(allow_patterns=["README.md", "model.pt", "config.json", "*.jit"])
 
     # Download the requested Autoregressive models
-    for size in args.model_sizes:
-        for model_type in args.model_types:
+    for size in model_sizes:
+        for model_type in model_types:
             suffix = f"-{model_type}"
             model_name = model_map[size] + suffix
             repo_id = f"{ORG_NAME}/{model_name}"
@@ -109,15 +109,11 @@ def main(args):
             local_dir_use_symlinks=False,
         )
 
-    if "Video2World" in args.model_types:
+    if "Video2World" in model_types:
         # Prompt Upsampler for Cosmos-1.0-Diffusion-Video2World models
         convert_pixtral_checkpoint(
-            checkpoint_dir=args.checkpoint_dir,
+            checkpoint_dir=checkpoint_dir,
             checkpoint_name="Pixtral-12B",
             vit_type="pixtral-12b-vit",
         )
 
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

text2world_hf.py

@@ -9,6 +9,7 @@ from .log import log
 from .misc import misc, Color, timer
 from .utils_io import read_prompts_from_file, save_video
 from .df_config_config import attrs  # this makes huggingface to download the file
+from .download_diffusion import main as download_diffusion
 
 
 # custom config class
@@ -133,5 +134,9 @@ class DiffusionText2World(PreTrainedModel):
         other_args = kwargs.copy()
         other_args.pop("config")
         config.update(other_args)
+        breakpoint()
+        model_sizes = ["7B",] if "7B" in config.diffusion_transformer_dir else ["14B",]
+        model_types = ["Text2World",]
+        download_diffusion(model_types, model_sizes, config.checkpoint_dir)
         model = cls(config)
         return model