o
    viV                     @   sh  U d Z ddlZddlZddlmZ ddlmZ ddlZddlmZm	Z	 ddl
mZ ddlmZmZ ddlmZ G d	d
 d
ejZdedefddZeeeef ZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd dejZi Zeeef e d< i Z!eeef e d< defd d!Z"e"ed"d#ed#d$d%d&d' e"ed(d)ed)d*d+d&d' e"ed,d-ed-d.d/d' e"ed0d1ed2d3ered4d5d0d6ned7d8d/d9 e"ed:d;ed<d3ered4d5d=d>ned?d@dAd>d9 e"edBdCedDd3ered4d5dEd>ned7d8d/d9 e"edFdGdHdIdJdKedLd3er6ed4dMdNd>ned?dOdPd>dQ ed4dRdSd6Z#e"edSdGdTdIdJdKedUd3e#dQ e"edVdWdXedYd3ed4dZd[d>d\ e"ed]dWd^ed_d3ed4d`d]d6d\ e"edadbdcdIdJdKeddd3ered4dMded>ned?dOdfd>dQ e"edgdhdidIdjdkdldmednd3ed?dodpd>dQ e"edqdhdrdIdjdkdldmedsd3ered4dMdtd>ned?dOdud>dQ e"edvdwdxdIdJdKedyd3ed?dzd{d>dQ e"ed|d}d~dIdJdKedd3er
ed4dMdd>nddQ e"eddddIdJdKedd3er&ed4dMdd>nddQ e"eddddddKedd3erBed4ddd>ned?ddd>dQ e"eddddIdJdKedd3erced4dMdd>nddQ e"edddddIdJdedd3ed4ddd>dQ e"edddddIdJdedd3ered4dMdd>nedddd>dQ e"edddddIdJdedd3ered4ddd>nedddd>dQ e"eddddIdJdKedd3ered4ddd>nedddd>dQ e"eddddIdJdKedd3ered4ddd>nedddd>dQ e"eddddIdJdKedd3ered4ddd>nedddd>dQ e"eddddIdJdKedd3er>ed4ddd>nedddd>dQ e"eddddIdJdKedd3er_ed4dMdd>nedddd>dQ e"eddddIdJdKedd3ered4dMdd>nedddd>dQ e"eddddIdJdKedd3ered4dMdd>nedddd>dQ e"eddddIdJdKedd3ered4dMdd>nedddd>dQ e"eddddIddkdldmedd3edddd>dQ e"eddddIdJdkdldmedd3ered4dMdd>nedddd>dQ dedefddZ$dedefddZ%ej&dedefddZ'ej&dedefddZ(dS )z!Database of released checkpoints.    N)cached_property)	Annotated)hf_hub_downloadsnapshot_download)override)EXPERIMENTAL_CHECKPOINTSINTERNAL)logc                   @   sV   e Zd ZU dZejdddZejedZ	ee
d< 	 defdd	Zedefd
dZdS )_CheckpointUriz%Config for checkpoint file/directory.forbidTextrafrozendefault_factorymetadatareturnc                 C   s   t d)Nz Download method not implemented.)NotImplementedErrorself r   W/data/cameron/vidgen/cosmos-policy/cosmos_policy/_src/imaginaire/utils/checkpoint_db.py	_download*   s   z_CheckpointUri._downloadc                 C   s   |   S )Return S3 URI or local path.)r   r   r   r   r   path-   s   z_CheckpointUri.pathN)__name__
__module____qualname____doc__pydantic
ConfigDictmodel_configFielddictr   __annotations__strr   r   r   r   r   r   r   r
      s   
 r
   urir   c                 C   s$   |  dstd|  d| dS )Ns3://zInvalid S3 URI: z. Must start with 's3://'/)
startswith
ValueErrorrstripr&   r   r   r   	is_s3_uri3   s   

r-   c                   @   s   e Zd ZU dZeed< dS )_CheckpointS3zConfig for checkpoint on S3.r&   N)r   r   r   r   S3Urir$   r   r   r   r   r.   <   s   
 r.   c                   @      e Zd ZdZdS )CheckpointFileS3z!Config for checkpoint file on S3.Nr   r   r   r   r   r   r   r   r1   C       r1   c                   @   r0   )CheckpointDirS3z&Config for checkpoint directory on S3.Nr2   r   r   r   r   r4   G   r3   r4   c                   @   s$   e Zd ZU dZeed< 	 eed< dS )_CheckpointHfz&Config for checkpoint on Hugging Face.
repositoryrevisionN)r   r   r   r   r%   r$   r   r   r   r   r5   K   s   
 r5   c                   @   s.   e Zd ZU dZeed< 	 edefddZdS )CheckpointFileHfz+Config for checkpoint file on Hugging Face.filenamer   c                 C   sL   t | jd| j| jd}td|  tdi |}tj	|s$J ||S ).Download checkpoint and return the local path.model)repo_id	repo_typer7   r9   z3Downloading checkpoint file from Hugging Face with Nr   )
r#   r6   r7   r9   r	   infor   osr   exists)r   download_kwargsr   r   r   r   r   Z   s   zCheckpointFileHf._downloadN)r   r   r   r   r%   r$   r   r   r   r   r   r   r8   T   s   
 r8   c                   @   s^   e Zd ZU dZdZeed< 	 dZeedf ed< 	 dZ	eedf ed< 	 e
defd	d
ZdS )CheckpointDirHfz0Config for checkpoint directory on Hugging Face. subdirectoryr   .includeexcluder   c                    s   i } j rt j |d< ndg|d<  jrt j|d<  jr* fdd| D }t jd jd|B }t	d|  t
d
i |} jrOtj| j}tj|sYJ ||S )r:   allow_patterns*ignore_patternsc                    s$   i | ]\}}| fd d|D qS )c                    s   g | ]
}t j j|qS r   )r?   r   joinrD   ).0xr   r   r   
<listcomp>   s    z8CheckpointDirHf._download.<locals>.<dictcomp>.<listcomp>r   )rK   keyvalr   r   r   
<dictcomp>   s   $ z-CheckpointDirHf._download.<locals>.<dictcomp>r;   )r<   r=   r7   z.Downloading checkpoint from Hugging Face with Nr   )rE   listrF   rD   itemsr#   r6   r7   r	   r>   r   r?   r   rJ   r@   )r   patternsrA   r   r   r   r   r   v   s   
zCheckpointDirHf._downloadN)r   r   r   r   rD   r%   r$   rE   tuplerF   r   r   r   r   r   r   rB   f   s   
 rB   c                   @   s   e Zd ZU dZejdddZeed< 	 eed< 	 ej	e
dZe
ed< 	 d	Zed	B ed
< 	 d	ZeeB d	B ed< 	 d	ZeeB d	B ed< 	 edefddZd	S )CheckpointConfigzConfig for checkpoint.r   Tr   uuidnamer   r   N
experiments3hfr   c                 C   s^   t r| jdur| jjS | jdu rtd| j d| j dtd| j d| j d | jj	S )r   NzCheckpoint (z#) is not available on Hugging Face.zDownloading checkpoint ))
r   rY   r&   rZ   r*   rW   rV   r	   r>   r   r   r   r   r   r      s   
zCheckpointConfig.path)r   r   r   r   r   r    r!   r%   r$   r"   r#   r   rX   rY   r1   r4   rZ   r8   rB   r   r   r   r   r   r   rU      s"   
 rU   _CHECKPOINTS_BY_UUID_CHECKPOINTS_BY_S3checkpoint_configc                 C   s`   | j tv rtd| j  d| t| j < | jd ur.| jj}|tv r(td| d| t|< d S d S )NCheckpoint UUID z already registered.Checkpoint S3 )rV   r]   r*   rY   r&   r^   )r_   r&   r   r   r   _register_checkpoint   s   


rb   z$4dbf13c6-1d30-4b02-99d6-75780dd8b744zgoogle-t5/t5-11bZ(90f37703b3334dfe9d2b009bfcbfbf1ac9d28ea3)ztf_model.h5)r6   r7   rF   )rV   rW   rZ   z$a2944743-cf8d-427e-a6fc-b3c03d807064zmeta-llama/Llama-Guard-3-8BZ(7327bd9f6efbbe6101dc6cc4736302b3cbb6e425)z
original/*z$9c7b7da4-2d95-45bb-9cb8-2eed954e9736znvidia/Cosmos-Guardrail1Z(d6d4bfa899a71454a700907664f3e88f503950cf)r6   r7   z$7219c6c7-f878-4137-bbdb-76842ea85e70zQwen/Qwen2.5-VL-7B-InstructzSs3://bucket/cosmos_reasoning1/pretrained/Qwen_tokenizer/Qwen/Qwen2.5-VL-7B-Instructr,   znvidia/Cosmos-ExperimentalZ(736a20b6cfbc38e42ba3f7e7d8efa1d886c20db1)r6   r7   rD   znvidia/Cosmos-Reason1-7BZ(3210bec0495fdc7a8d3dbb8d58da5711eab4b423)rV   rW   rY   rZ   z$685afcaa-4de2-42fe-b7b9-69f7a2dee4d8z
Wan2.1/vaezQs3://bucket/cosmos_diffusion_v2/pretrain_weights/tokenizer/wan2pt1/Wan2.1_VAE.pthz(685afcaa-4de2-42fe-b7b9-69f7a2dee4d8.pth)r6   r7   r9   znvidia/Cosmos-Predict2.5-2BZ(6787e176dce74a101d922174a95dba29fa5f0c55ztokenizer.pthz$cb3e3ffa-7b08-4c34-822d-61c7aa31a14fznvidia/Cosmos-Reason1.1-7Bzs3://bucket/cosmos_reasoning1/sft_exp700/sft_exp721-1_qwen7b_tl_721_5vs5_s3_balanced_n32_resume_16k/checkpoints/iter_000016000/modelz-cb3e3ffa-7b08-4c34-822d-61c7aa31a14f/model.ptz$d20b7120-df3e-4911-919d-db6e08bad31cz,nvidia/Cosmos-Predict2.5-2B/base/pre-trainedzStage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweighted_1_1_rectified_flow_only_resume2720p   )
resolutionfpszs3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/Stage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweighted_1_1_rectified_flow_only_resume2/checkpoints/iter_000023000/modelZ(9a02ed8daa8c6c7718ac09da06488bfd1d363cb6z6d20b7120-df3e-4911-919d-db6e08bad31c/model_ema_bf16.ptZ(15a82a2ec231bc318692aa0456a36537c806e7d4zAbase/pre-trained/d20b7120-df3e-4911-919d-db6e08bad31c_ema_bf16.pt)rV   rW   rX   r   rY   rZ   Z(eda2f0ca1db6281c9a960908bb6bf14607a0fea0z$308eb96c-c4c0-4a06-9cc1-103a43beff28zgStage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweightedzs3://bucket/cosmos_diffusion_v2/official_runs_text2world/Stage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweighted/checkpoints/iter_000010000/z$7bbc8d06-2bc9-448d-94ee-b48b4ab7189cz'nvidia/Cosmos-Predict2.5-2B/interactiveZNcosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_sf_warmupzs3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000014000/modelZ(2b5e9a99b58d5a61259ca99962c4c74127481006z67bbc8d06-2bc9-448d-94ee-b48b4ab7189c/model_ema_bf16.pt)rV   rW   rX   rY   rZ   z$bedc35da-1a54-4144-83db-6072c29b0fd9Z,cosmos_predict2p5_2B_action_gr00t_gr1_warmupzfs3://bucket/cosmos_predict2_action_conditioned/interactive_warmup/gr1/checkpoints/iter_000020000/modelZ(ded876a5b2e19aef64cd9d1100c03e5b05cf2f9cz$81edfebe-bd6a-4039-8c1d-737df1a790bfz-nvidia/Cosmos-Predict2.5-2B/base/post-trainedzAStage-c_pt_4-Index-2-Size-2B-Res-720-Fps-16-Note-rf_with_edm_ckptzs3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/Stage-c_GRPO-reason_embeddings-Index-26-Size-2B-Res-720-Fps-16-posttrain_data-HQ_V7_RF_MERGE_LOCAL_ag_every2_guidance0_scorekeyoverall_reward_databeta0.01_mincon0/checkpoints/iter_000000288/modelz681edfebe-bd6a-4039-8c1d-737df1a790bf/model_ema_bf16.ptzBbase/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.ptz$524af350-2e43-496c-8590-3646ae1325daz*nvidia/Cosmos-Predict2.5-2B/auto/multiviewZbuttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0         )re   rf   viewsframeszs3://bucket/cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0-0/checkpoints/iter_000012000/modelZ(865baf084d4c9e850eac59a021277d5a9b9e8b63z?auto/multiview/524af350-2e43-496c-8590-3646ae1325da_ema_bf16.ptz$6b9d7548-33bb-4517-b5e8-60caf47edba7Zbuttercup_predict2p5_2b_7views_res720p_fps30_t8_from48kfps30mv_condprobs0442_joint_alpamayo1capnoviewprefix_allcapsviewprefix_29frames_nofpszs3://bucket/cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_7views_res720p_fps30_t8_from48kfps30mv_condprobs0442_joint_alpamayo1capnoviewprefix_allcapsviewprefix_29frames_nofps-0/checkpoints/iter_000005000/modelz66b9d7548-33bb-4517-b5e8-60caf47edba7/model_ema_bf16.ptz?auto/multiview/6b9d7548-33bb-4517-b5e8-60caf47edba7_ema_bf16.ptz$f740321e-2cd6-4370-bbfe-545f4eca2065z2nvidia/Cosmos-Predict2.5-2B/robot/multiview-agibotZKmulticamera_video2video_rectified_flow_2b_res_720_fps16_s3_agibot_frameinitzs3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/multicamera_video2video_rectified_flow_2b_res_720_fps16_s3_agibot_frameinit/checkpoints/iter_000016500Z(fbe72c18d152053029a19db3b211cf78671ad422z6f740321e-2cd6-4370-bbfe-545f4eca2065/model_ema_bf16.ptz$0e8177cc-0db5-4cfd-a8a4-b820c772f4fcz5nvidia/Cosmos-Transfer2.5-2B/general/multiview-cameraZJmulticamera_video2video_rectified_flow_2b_res_720_fps16_s3_multicam_syncamzs3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/multicamera_video2video_rectified_flow_2b_res_720_fps16_s3_multicam_syncam/checkpoints/iter_000002000/modelz60e8177cc-0db5-4cfd-a8a4-b820c772f4fc/model_ema_bf16.ptz$7f6b99b7-7fac-4e74-8dbe-a394cb56ef99z:nvidia/Cosmos-Transfer2.5-2B/robot/multiview-agibot-cameraZAmulticamera_video2video_rectified_flow_2b_res_720_fps16_s3_agibotzs3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/multicamera_video2video_rectified_flow_2b_res_720_fps16_s3_agibot/checkpoints/iter_000003000/modelz67f6b99b7-7fac-4e74-8dbe-a394cb56ef99/model_ema_bf16.ptz$38c6c645-7d41-4560-8eeb-6f4ddc0e6574z-nvidia/Cosmos-Predict2.5-2B/robot/action-condZ_cosmos_predict2p5_2B_reason_embeddings_action_conditioned_rectified_flow_bridge_13frame_256x320Z360p   zs3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_reason_embeddings_action_conditioned_rectified_flow_bridge_13frame_256x320/checkpoints/iter_000016000/modelmainz638c6c645-7d41-4560-8eeb-6f4ddc0e6574/model_ema_bf16.ptzBrobot/action-cond/38c6c645-7d41-4560-8eeb-6f4ddc0e6574_ema_bf16.ptz$24a3b7b8-6a3d-432d-b7d1-5d30b9229465z'nvidia/Cosmos-Predict2.5-2B/transfer2.5zStage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweighted_1_1_rectified_flow_onlyzs3://bucket/cosmos_diffusion_v2/official_runs_text2world/Stage-c_pt_4-reason_embeddings-v1p1-Index-26-Size-2B-Res-720-Fps-16-Note-T2V_high_sigma_loss_reweighted_1_1_rectified_flow_only/checkpoints/iter_000037000/modelz624a3b7b8-6a3d-432d-b7d1-5d30b9229465/model_ema_bf16.ptz$575edf0f-d973-4c74-b52c-69929a08d0a5z*nvidia/Cosmos-Predict2.5-2B/base/distilledZ<dmd2_trigflow_distill_cosmos_predict2_2B_bidirectional_TnI2VZ2B)sizere   rf   zs3://bucket/cosmos_predict2_distill/predict2_distill/dmd2_trigflow_distill_cosmos_predict2_2B_bidirectional/checkpoints/iter_000007500/modelZ(cb56c64d7e5bb20a50b1e39b4429b967522e91d4z6575edf0f-d973-4c74-b52c-69929a08d0a5/model_ema_bf16.ptz$54937b8c-29de-4f04-862c-e67b04ec41e8z-nvidia/Cosmos-Predict2.5-14B/base/pre-trainedz{Stage-c_pt_4-reason_embeddings-v1p1-Index-43-Size-14B-Res-720-Fps-16_resume_from_reason1p1_rectified_flow_shift5_high_sigmaZ14Bzs3://bucket/cosmos_diffusion_v2/official_runs_text2world/Stage-c_pt_4-reason_embeddings-v1p1-Index-43-Size-14B-Res-720-Fps-16_resume_from_reason1p1_rectified_flow_shift5_high_sigma/checkpoints/iter_000012500/modelz654937b8c-29de-4f04-862c-e67b04ec41e8/model_ema_bf16.ptznvidia/Cosmos-Predict2.5-14BZ(03eb354f35eae0d6e0c1be3c9f94d8551e125570zAbase/pre-trained/54937b8c-29de-4f04-862c-e67b04ec41e8_ema_bf16.ptz$e21d2a49-4747-44c8-ba44-9f6f9243715fz.nvidia/Cosmos-Predict2.5-14B/base/post-traineda  s3://bucket/cosmos_diffusion_v2/official_runs_vid2vid/Stage-c_GRPO-reason_embeddings-Index-26-Size-14B-Res-720-Fps-16-posttrain_data-HQ_V7_RF_MERGE_GENERAL_steps20_every2_lr3e-6_guidance0_scorekeyoverall_reward_databeta0.01_mincon0/checkpoints/iter_000000128/modelZ(9e46ea0945ac591f9d4abe810ecd78cde112fb82z6e21d2a49-4747-44c8-ba44-9f6f9243715f/model_ema_bf16.ptZ(2bc4ca5ba5a20b9858a7ddb856bc82d70b030fbezBbase/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.ptz$61f5694b-0ad5-4ecd-8ad7-c8545627d125z)nvidia/Cosmos-Transfer2.5-2B/general/edgezqedge_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv3p1_20250714_64N_rectified_flow_refimdrop0pt5zs3://bucket/cosmos_transfer2/vid2vid_2B_control/edge_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv3p1_20250714_64N_rectified_flow_refimdrop0pt5/checkpoints/iter_000032000/modelZ(61c640b2f4092cb0868c9d9941fa505a750ccd4dz661f5694b-0ad5-4ecd-8ad7-c8545627d125/model_ema_bf16.ptznvidia/Cosmos-Transfer2.5-2BZ(b67b64abda3801a9aceddbff2bdb86126c06db74z=general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.ptz$626e6618-bfcd-4d9a-a077-1409e2ce353fz*nvidia/Cosmos-Transfer2.5-2B/general/depthzrdepth_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv4p1_20250823_64N_rectified_flow_refimdrop0pt5zs3://bucket/cosmos_transfer2/vid2vid_2B_control/depth_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv4p1_20250823_64N_rectified_flow_refimdrop0pt5/checkpoints/iter_000044000/modelz6626e6618-bfcd-4d9a-a077-1409e2ce353f/model_ema_bf16.ptZ(dea7737ca29dd8d9086413c6dc5724b8250a0bb4z>general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.ptz$ba2f44f2-c726-4fe7-949f-597069d9b91cz)nvidia/Cosmos-Transfer2.5-2B/general/blurz}vis_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv3p1_20250714_64N_rectified_flow_refimdrop0pt5_filterb3g5m2zs3://bucket/cosmos_transfer2/vid2vid_2B_control/vis_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv3p1_20250714_64N_rectified_flow_refimdrop0pt5_filterb3g5m2/checkpoints/iter_000036000/Z(8ecf7ad717c10b9f796bab731eba311e480fcf58z6ba2f44f2-c726-4fe7-949f-597069d9b91c/model_ema_bf16.ptZ(eb5325b77d358944da58a690157dd2b8071bbf85z=general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.ptz$5136ef49-6d8d-42e8-8abf-7dac722a304az(nvidia/Cosmos-Transfer2.5-2B/general/segzpseg_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv4p2_20250823_64N_rectified_flow_refimdrop0pt5zs3://bucket/cosmos_transfer2/vid2vid_2B_control/seg_720p_t24or1_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_uniform_hqv4p2_20250823_64N_rectified_flow_refimdrop0pt5/checkpoints/iter_000043000/z65136ef49-6d8d-42e8-8abf-7dac722a304a/model_ema_bf16.ptZ(23057a4167b89de89a4a397fdbf3887994d115ebz<general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.ptz$ecd0ba00-d598-4f94-aa09-e8627899c431zmedge_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv3p1_20250714_64N_rectified_flow_mock_datazs3://bucket/cosmos_transfer2/vid2vid_2B_control/edge_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv3p1_20250714_64N_rectified_flow/checkpoints/iter_000029000/modelz6ecd0ba00-d598-4f94-aa09-e8627899c431/model_ema_bf16.ptZ(bd963eabcfc2d61dc4ea365cacf41d45ac480aa5z=general/edge/ecd0ba00-d598-4f94-aa09-e8627899c431_ema_bf16.ptz$fcab44fe-6fe7-492e-b9c6-67ef8c1a52abzbseg_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv4p2_20250823_64N_rectified_flowzs3://bucket/cosmos_transfer2/vid2vid_2B_control/seg_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv4p2_20250823_64N_rectified_flow/checkpoints/iter_000031000/modelz6fcab44fe-6fe7-492e-b9c6-67ef8c1a52ab/model_ema_bf16.ptz<general/seg/fcab44fe-6fe7-492e-b9c6-67ef8c1a52ab_ema_bf16.ptz$20d9fd0b-af4c-4cca-ad0b-f9b45f0805f1zbvis_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv3p1_20250714_64N_rectified_flowzs3://bucket/cosmos_transfer2/vid2vid_2B_control/vis_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv3p1_20250714_64N_rectified_flow/checkpoints/iter_000043000/modelz620d9fd0b-af4c-4cca-ad0b-f9b45f0805f1/model_ema_bf16.ptz=general/blur/20d9fd0b-af4c-4cca-ad0b-f9b45f0805f1_ema_bf16.ptz$0f214f66-ae98-43cf-ab25-d65d09a7e68fzddepth_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv4p1_20250823_64N_rectified_flowzs3://bucket/cosmos_transfer2/vid2vid_2B_control/depth_720p_t24_spaced_layer4_cr1pt1_sdev2_lowsigma0.05_nonuniform_hqv4p1_20250823_64N_rectified_flow/checkpoints/iter_000028000/modelz60f214f66-ae98-43cf-ab25-d65d09a7e68f/model_ema_bf16.ptz>general/depth/0f214f66-ae98-43cf-ab25-d65d09a7e68f_ema_bf16.ptz$4ecc66e9-df19-4aed-9802-0d11e057287az+nvidia/Cosmos-Transfer2.5-2B/auto/multiviewZbuttercup_transfer2p5_2b_mv_7views_res720p_fps10_t8_fromfinetuned12knofpsuniform_mads720pmulticaps29frames_world_scenario_nofps_uniform
   zs3://bucket/cosmos_transfer2_multiview/cosmos2_mv/buttercup_transfer2p5_2b_mv_7views_res720p_fps10_t8_fromfinetuned12knofpsuniform_mads720pmulticaps29frames_world_scenario_nofps_uniform-0/checkpoints/iter_000006500/modelZ(00c591edab119e8a6ca06e6e091351a04ce0ecc9z?auto/multiview/4ecc66e9-df19-4aed-9802-0d11e057287a_ema_bf16.ptz$b5ab002d-a120-4fbf-a7f9-04af8615710bZzbuttercup_transfer2p5_2b_mv_7views_res720p_fps10_t8_frombase5knofps_mads720pmulticaps29frames_world_scenario_resumefrom21kzs3://bucket/cosmos_transfer2_multiview/cosmos2_mv/buttercup_transfer2p5_2b_mv_7views_res720p_fps10_t8_frombase5knofps_mads720pmulticaps29frames_world_scenario_resumefrom21k-0/checkpoints/iter_000010000/modelz6b5ab002d-a120-4fbf-a7f9-04af8615710b/model_ema_bf16.ptz?auto/multiview/b5ab002d-a120-4fbf-a7f9-04af8615710b_ema_bf16.ptcheckpoint_uuidc                 C   s    | t vrtd|  dt |  S )z"Return checkpoint config for UUID.r`    not found.)r]   r*   )rp   r   r   r   get_checkpoint_by_uuid  s   rr   checkpoint_s3c                 C   s*   |  d} | tvrtd|  dt|  S )z$Return checkpoint config for S3 URI.r(   ra   rq   )r+   r^   r*   )rs   r   r   r   get_checkpoint_by_s3  s   
rt   checkpoint_hfc                 C   s   |  dsJ d|  | dd }|d}t|dk r%td|  dd|dd	 }d|d	d }td
| d|  t|d|d}tj	
|sSJ ||S )z;Download checkpoint from HuggingFace and return local path.hf://zNot a HuggingFace URI:    Nr(      z Invalid HuggingFace URI format: z1. Expected format: hf://org/repo/path/to/file.pth   z)Downloading checkpoint from HuggingFace: r;   )r<   r=   r9   )r)   splitlenr*   rJ   r	   r>   r   r?   r   r@   )ru   hf_pathpartsr<   r9   r   r   r   r   get_checkpoint_by_hf  s"   

r~   checkpoint_uric                 C   sX   t r| S | d} | drt| jS | drt| S tj| s*td|  d| S )zReturn checkpoint path for S3 URI, HuggingFace URI, or local path.

    Supports:
    - S3 URIs: s3://bucket/path/to/checkpoint
    - HuggingFace URIs: hf://org/repo/path/to/file.pth
    - Local paths: /path/to/checkpoint
    r(   r'   rv   zCheckpoint path z does not exist.)	r   r+   r)   rt   r   r~   r?   r@   r*   )r   r   r   r   get_checkpoint_path  s   	



r   ))r   	functoolsr?   r   typingr   r   huggingface_hubr   r   typing_extensionsr   #cosmos_policy._src.imaginaire.flagsr   r   #cosmos_policy._src.imaginaire.utilsr	   	BaseModelr
   r%   r-   AfterValidatorr/   r.   r1   r4   r5   r8   rB   rU   r]   r#   r$   r^   rb   ru   rr   rt   	lru_cacher~   r   r   r   r   r   <module>   s>  	%$