o vi @s<dZddlmZddlmZd dededBdefdd ZdS) zS Checkpoint utilities for Cosmos Policy, including HuggingFace checkpoint loading. )hf_hub_download)logNcheckpoint_path cache_dirreturnc Cs|drI|tdd}|dd}t|dkr!td|d|\}}}|d|}td|d|t|||d }td ||S|S) a Resolve checkpoint path, downloading from HuggingFace if needed. Supports: - Local paths: /path/to/checkpoint.pth - S3 paths: s3://bucket/path/to/checkpoint.pth - HuggingFace paths: hf://org/repo-name/path/to/file.pth Args: checkpoint_path: Path to checkpoint file cache_dir: Optional cache directory for HuggingFace downloads Returns: Resolved local path to checkpoint Examples: >>> # HuggingFace download >>> path = resolve_checkpoint_path("hf://nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth") >>> # S3 or local path (passed through unchanged) >>> path = resolve_checkpoint_path("s3://bucket/checkpoint.pth") >>> path = resolve_checkpoint_path("/local/path/checkpoint.pth") zhf://N/z!Invalid HuggingFace path format: z2. Expected format: hf://org/repo-name/path/to/filez)Downloading checkpoint from HuggingFace: )repo_idfilenamerzDownloaded checkpoint to: ) startswithlensplit ValueErrorrinforsuccess) rrhf_pathpartsorg repo_namer r local_pathrJ/data/cameron/vidgen/cosmos-policy/cosmos_policy/utils/checkpoint_utils.pyresolve_checkpoint_paths$     r)N)__doc__huggingface_hubr#cosmos_policy._src.imaginaire.utilsrstrrrrrrs