add download model and how to use

README.md

@@ -1,3 +1,38 @@
+## How to Use
+
+'''python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "NeverMore0123/AutoregressiveVideo2WorldGeneration", 
+    cache_dir="./cache", 
+    trust_remote_code=True,
+    
+    input_type = "text_and_image",
+    num_input_frames = 1,
+    prompt = "A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." ,
+    input_image_or_video_path = "AutoregressiveVideo2WorldGeneration/cosmos1/models/autoregressive/assets/v1p0/input.jpg",
+    video_save_name = "diffusion_decoder_image_output",
+    ar_model_dir = "Cosmos-1.0-Autoregressive-5B-Video2World",
+    
+    # input_type = "text_and_video",
+    # num_input_frames = 9,
+    # prompt = "A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." ,
+    # input_image_or_video_path = "AutoregressiveVideo2WorldGeneration/cosmos1/models/autoregressive/assets/v1p0/input.mp4",
+    # video_save_name = "diffusion_decoder_video_output",
+    
+    # turn on offloading on a low GPU memory machine:
+    disable_diffusion_decoder=False,
+    offload_guardrail_models=True,
+    offload_diffusion_decoder=True,
+    offload_network=True,
+    offload_tokenizer=True,
+    offload_text_encoder_model=True,
+)
+
+model()
+'''
+
 
 ![Cosmos Logo](assets/cosmos-logo.png)
 

download_autoregressive.py

@@ -18,34 +18,7 @@ from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Download NVIDIA Cosmos-1.0 Autoregressive models from Hugging Face")
-    parser.add_argument(
-        "--model_sizes",
-        nargs="*",
-        default=[
-            "4B",
-            "5B",
-        ],  # Download all by default
-        choices=["4B", "5B", "12B", "13B"],
-        help="Which model sizes to download. Possible values: 4B, 5B, 12B, 13B.",
-    )
-    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(args):
+def main(model_sizes, checkpoint_dir="checkpoints"):
     ORG_NAME = "nvidia"
 
     # Mapping from size argument to Hugging Face repository name
@@ -65,13 +38,13 @@ def main(args):
     ]
 
     # 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 size in model_sizes:
         model_name = model_map[size]
         repo_id = f"{ORG_NAME}/{model_name}"
         local_dir = checkpoints_dir.joinpath(model_name)
@@ -99,7 +72,3 @@ def main(args):
             local_dir_use_symlinks=False,
         )
 
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

video2world_hf.py

@@ -24,7 +24,7 @@ from .ar_utils_inference import load_vision_input, validate_args
 from .log import log
 from .io import read_prompts_from_file
 
-# from download_autoregressive import main as download_autoregressive
+from .download_autoregressive import main as download_autoregressive
 from transformers import PreTrainedModel, PretrainedConfig
 
 
@@ -142,8 +142,7 @@ class ARVideo2World(PreTrainedModel):
         other_args = kwargs.copy()
         other_args.pop("config")
         config.update(other_args)
-        # model_sizes = ["5B",] if "5B" in config.ar_model_dir else ["13B",]
-        # model_types = ["Video2World",]
-        # download_autoregressive(model_types, model_sizes, config.checkpoint_dir)
+        model_sizes = ["5B",] if "5B" in config.ar_model_dir else ["13B",]
+        download_autoregressive(model_sizes, config.checkpoint_dir)
         model = cls(config)
         return model