should work now

download_diffusion.py

@@ -21,42 +21,6 @@ from huggingface_hub import snapshot_download
 from .convert_pixtral_ckpt import convert_pixtral_checkpoint
 
 
-def parse_args():
-    parser = argparse.ArgumentParser(description="Download NVIDIA Cosmos-1.0 Diffusion models from Hugging Face")
-    parser.add_argument(
-        "--model_sizes",
-        nargs="*",
-        default=[
-            "7B",
-            "14B",
-        ],  # Download all by default
-        choices=["7B", "14B"],
-        help="Which model sizes to download. Possible values: 7B, 14B",
-    )
-    parser.add_argument(
-        "--model_types",
-        nargs="*",
-        default=[
-            "Text2World",
-            "Video2World",
-        ],  # Download all by default
-        choices=["Text2World", "Video2World"],
-        help="Which model types to download. Possible values: Text2World, Video2World",
-    )
-    parser.add_argument(
-        "--cosmos_version",
-        type=str,
-        default="1.0",
-        choices=["1.0"],
-        help="Which version of Cosmos to download. Only 1.0 is available at the moment.",
-    )
-    parser.add_argument(
-        "--checkpoint_dir", type=str, default="checkpoints", help="Directory to save the downloaded checkpoints."
-    )
-    args = parser.parse_args()
-    return args
-
-
 def main(model_types, model_sizes, checkpoint_dir="checkpoints"):
     ORG_NAME = "nvidia"
 

text2world_hf.py

@@ -48,10 +48,10 @@ class DiffusionText2World(PreTrainedModel):
 
     def __init__(self, config=DiffusionText2WorldConfig()):
         super().__init__(config)
-        torch.enable_grad(False)   # TODO: do we need this?
+        torch.enable_grad(False) 
         self.config = config
         inference_type = "text2world"
-        config.prompt = 1          # TODO: this is to hack args validation, maybe find a better way 
+        config.prompt = 1          # this is to hack args validation, maybe find a better way 
         validate_args(config, inference_type) 
         del config.prompt
         self.pipeline = DiffusionText2WorldGenerationPipeline(