(args, transformer, vae)
| 755 | |
| 756 | |
| 757 | def save_pipeline_cosmos_1_0(args, transformer, vae): |
| 758 | text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16) |
| 759 | tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path) |
| 760 | # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly. |
| 761 | # So, the sigma_min values that is used is the default value of 0.002. |
| 762 | scheduler = EDMEulerScheduler( |
| 763 | sigma_min=0.002, |
| 764 | sigma_max=80, |
| 765 | sigma_data=0.5, |
| 766 | sigma_schedule="karras", |
| 767 | num_train_timesteps=1000, |
| 768 | prediction_type="epsilon", |
| 769 | rho=7.0, |
| 770 | final_sigmas_type="sigma_min", |
| 771 | ) |
| 772 | |
| 773 | pipe_cls = CosmosTextToWorldPipeline if "Text2World" in args.transformer_type else CosmosVideoToWorldPipeline |
| 774 | pipe = pipe_cls( |
| 775 | text_encoder=text_encoder, |
| 776 | tokenizer=tokenizer, |
| 777 | transformer=transformer, |
| 778 | vae=vae, |
| 779 | scheduler=scheduler, |
| 780 | safety_checker=lambda *args, **kwargs: None, |
| 781 | ) |
| 782 | pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") |
| 783 | |
| 784 | |
| 785 | def save_pipeline_cosmos_2_0(args, transformer, vae): |
no test coverage detected
searching dependent graphs…